tvm.te

张量表达式语言的命名空间

功能：

`any`(*args[, span])	创建一个新表达式，表示参数中所有条件的并集
`all`(*args[, span])	创建一个新表达式，表示所有条件的交集
`min_value`(dtype[, span])	数据类型的最小值
`max_value`(dtype[, span])	数据类型的最大值
`trace`(args[, trace_action])	在运行时追踪张量数据。
`exp`(x)	对输入x取指数。
`erf`(x)	对输入x取高斯误差函数。
`tanh`(x)	对输入x取双曲正切tanh。
`sigmoid`(x)	快速获取sigmoid值的函数
`log`(x)	对输入x取对数。
`tan`(x)	计算输入x的正切值。
`cos`(x)	计算输入x的余弦值。
`sin`(x)	计算输入x的正弦值。
`sqrt`(x)	对输入x取平方根。
`rsqrt`(x)	取输入x的平方根的倒数。
`floor`(x[, span])	对浮点输入x取地板值。
`ceil`(x[, span])	对浮点输入x取上界。
`sinh`(x)	计算输入x的双曲正弦值。
`cosh`(x)	计算输入x的双曲余弦值。
`log2`(x)	对输入x取以2为底的对数。
`log10`(x)	对输入x取以10为底的对数。
`asin`(x)	取输入x的asin值。
`asinh`(x)	对输入x取反双曲正弦函数asinh。
`acos`(x)	计算输入x的反余弦值。
`acosh`(x)	计算输入x的反余弦值。
`atan`(x)	计算输入x的反正切值。
`atanh`(x)	对输入x取反双曲正切。
`trunc`(x[, span])	获取输入的截断值。
`abs`(x[, span])	获取输入元素的逐元素绝对值。
`round`(x[, span])	将数组中的元素四舍五入到最接近的整数。
`nearbyint`(x[, span])	将数组中的元素四舍五入到最接近的整数。
`power`(x, y[, span])	x的y次方
`popcount`(x)	计算输入x中置位（set bit）的数量。
`fmod`(x, y)	返回x除以y的余数，其符号与x相同。
`if_then_else`(cond, t, f[, span])	条件选择表达式。
`isnan`(x[, span])	检查输入值是否为NaN。
`isfinite`(x[, span])	检查输入值是否为有限值。
`isinf`(x[, span])	检查输入值是否为无穷大。
`div`(a, b[, span])	按照C/C++语义计算a / b。
`indexdiv`(a, b[, span])	计算 floor(a / b)，其中 a 和 b 为非负数。
`indexmod`(a, b[, span])	计算indexdiv的余数。
`truncdiv`(a, b[, span])	计算两个表达式的截断除法。
`truncmod`(a, b[, span])	计算两个表达式的截断模。
`floordiv`(a, b[, span])	计算两个表达式的向下取整除。
`floormod`(a, b[, span])	计算两个表达式的取模运算结果。
`comm_reducer`(fcombine, fidentity[, name])	为归约操作创建一个可交换的归约器。
`min`(expr, axis[, where, init])	在轴上创建一个最小表达式。
`max`(expr, axis[, where, init])	创建一个沿轴的最大值表达式。
`sum`(expr, axis[, where, init])	创建一个沿轴求和的表达式。
`add`(lhs, rhs[, span])	通用加法运算符。
`subtract`(lhs, rhs[, span])	通用减法运算符。
`multiply`(lhs, rhs[, span])	通用乘法运算符。
`tag_scope`(tag)	操作符标签的作用范围。
`placeholder`(shape[, dtype, name])	构造一个空的张量对象。
`compute`(shape, fcompute[, name, tag, attrs, ...])	通过计算形状域构造一个新的张量。
`scan`(init, update, state_placeholder[, ...])	通过沿轴扫描构建新的张量。
`extern`(shape, inputs, fcompute[, name, ...])	通过外部函数计算多个张量。
`var`([name, dtype, span])	创建一个具有指定名称和数据类型的新变量
`size_var`([name, dtype, span])	创建一个新变量表示张量形状大小，该大小必须为非负数。
`const`(value[, dtype, span])	创建一个具有指定值和数据类型的新常量
`thread_axis`([dom, tag, name, span])	创建一个新的IterVar来表示线程索引。
`reduce_axis`(dom[, 名称, 线程标签, 范围])	创建一个新的归约迭代变量。
`create_prim_func`(ops[, index_dtype_override])	从张量表达式创建TensorIR PrimFunc
`extern_primfunc`(input_tensors, primfunc, ...)	通过可调度的TIR PrimFunc计算张量

类：

`TensorSlice`(tensor, indices)	用于支持张量切片语法的辅助数据结构。
`Tensor`	张量对象，要构造它，请参见function.Tensor
`PlaceholderOp`	占位符操作。
`ComputeOp`	标量运算。
`ScanOp`	扫描操作。
`ExternOp`	外部操作。

tvm.te.any(*args, span=None)

创建一个新表达式，表示参数中所有条件的并集

Parameters:

args (list) – 符号布尔表达式列表
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

expr – 表达式

Return type:

表达式

tvm.tir.any()的别名

tvm.te.all(*args, span=None)

Create a new expression of the intersection of all conditions in the: 参数

Parameters:

args (list) – 符号化布尔表达式列表
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

expr – 表达式

Return type:

表达式

tvm.tir.all()的别名

tvm.te.min_value(dtype, span=None)

数据类型的最小值

Parameters:

dtype (str) – 数据类型。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

value – dtype的最小值。

Return type:

tvm.Expr

tvm.tir.min_value()的别名

tvm.te.max_value(dtype: str, span: Span | None = None) → 任意

数据类型的最大值

Parameters:

dtype (str) – 数据类型。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

value – dtype的最大值。

Return type:

tvm.Expr

tvm.tir.max_value()的别名

tvm.te.trace(args, trace_action='tvm.default_trace_action')

在运行时追踪张量数据。

trace函数允许在运行时追踪特定的张量。追踪值应作为最后一个参数传入。需要指定追踪动作，默认情况下会使用tvm.default_trace_action。

Parameters:

args (list of Expr or Buffers.) - 位置参数。
trace_action (str.) – 追踪动作的名称。

Returns:

call – 调用表达式。

Return type:

PrimExpr

另请参阅

tvm.tir.call_packed: 创建打包函数。

tvm.tir.trace()的别名

tvm.te.exp(x)

对输入x取指数。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.exp()的别名

tvm.te.erf(x)

对输入x取高斯误差函数。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.erf()的别名

tvm.te.tanh(x)

对输入x取双曲正切tanh。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.tanh()的别名

tvm.te.sigmoid(x)

快速获取sigmoid值的函数

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.sigmoid() 的别名

tvm.te.log(x)

对输入x取对数。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.log()的别名

tvm.te.tan(x)

计算输入x的正切值。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.tan()的别名

tvm.te.cos(x)

计算输入x的余弦值。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.cos()的别名

tvm.te.sin(x)

计算输入x的正弦值。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.sin()的别名

tvm.te.sqrt(x)

对输入x取平方根。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.sqrt()的别名

tvm.te.rsqrt(x)

取输入x的平方根的倒数。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.rsqrt()的别名

tvm.te.floor(x: PrimExprWithOp, span=None)

对浮点输入x取地板值。

Parameters:

x (PrimExpr) – 输入参数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

y – 结果。

Return type:

PrimExpr

tvm.tir.floor()的别名

tvm.te.ceil(x, span=None)

对浮点输入x取上界。

Parameters:

x (PrimExpr) – 输入参数。
span (可选参数[Span]) – 该操作符在源代码中的位置。

Returns:

y – 结果。

Return type:

PrimExpr

tvm.tir.ceil()的别名

tvm.te.sinh(x)

计算输入x的双曲正弦值。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.sinh() 的别名

tvm.te.cosh(x)

计算输入x的双曲余弦值。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.cosh()的别名

tvm.te.log2(x)

对输入x取以2为底的对数。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.log2()的别名

tvm.te.log10(x)

对输入x取以10为底的对数。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.log10()的别名

tvm.te.asin(x)

取输入x的asin值。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.asin()的别名

tvm.te.asinh(x)

对输入x取反双曲正弦函数asinh。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.asinh()的别名

tvm.te.acos(x)

计算输入x的反余弦值。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.acos()的别名

tvm.te.acosh(x)

计算输入x的反余弦值。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.acosh()的别名

tvm.te.atan(x)

计算输入x的反正切值。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.atan()的别名

tvm.te.atanh(x)

对输入x取反双曲正切。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.atanh()的别名

tvm.te.trunc(x, span=None)

获取输入的截断值。

标量x的截断值是最接近零的整数i，且i比x更接近零。

Parameters:

x (PrimExpr) – 输入参数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

y – 结果。

Return type:

PrimExpr

tvm.tir.trunc()的别名

tvm.te.abs(x, span=None)

获取输入元素的逐元素绝对值。

Parameters:

x (PrimExpr) – 输入参数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

y – 结果。

Return type:

PrimExpr

tvm.tir.abs()的别名

tvm.te.round(x, span=None)

将数组中的元素四舍五入到最接近的整数。

Parameters:

x (PrimExpr) – 输入参数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

y – 结果。

Return type:

PrimExpr

tvm.tir.round()的别名

tvm.te.nearbyint(x, span=None)

将数组元素四舍五入到最接近的整数。此内置函数使用llvm.nearbyint而非llvm.round，虽然速度更快但结果会与te.round不同。值得注意的是nearbyint会根据舍入模式进行舍入，而te.round（llvm.round）会忽略这一点。两者差异详见： https://en.cppreference.com/w/cpp/numeric/math/round https://en.cppreference.com/w/cpp/numeric/math/nearbyint

Parameters:

x (PrimExpr) – 输入参数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

y – 结果。

Return type:

PrimExpr

tvm.tir.nearbyint()的别名

tvm.te.power(x, y, span=None)

x的y次方

Parameters:

x (PrimExpr) – 输入参数。
y (PrimExpr) – 指数
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

z – 结果。

Return type:

PrimExpr

tvm.tir.power()的别名

tvm.te.popcount(x)

计算输入x中置位（set bit）的数量。

Parameters:: x (PrimExpr) – 输入参数。
Returns:: y – 结果。
Return type:: PrimExpr

tvm.tir.popcount()的别名

tvm.te.fmod(x, y)

返回x除以y的余数，其符号与x相同。

Parameters:

x (PrimExpr) – 输入参数。
y (PrimExpr) – 输入参数。

Returns:

z – 结果。

Return type:

PrimExpr

tvm.tir.fmod()的别名

tvm.te.if_then_else(cond, t, f, span=None)

条件选择表达式。

Parameters:

cond (PrimExpr) – 条件表达式
t (PrimExpr) – 如果条件cond为真时的结果表达式。
f (PrimExpr) – 如果条件cond为假时的结果表达式。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

result – 条件表达式的结果。

Return type:

Node

注意

与Select不同，if_then_else不会执行不满足条件的分支。您可以用它来防止越界访问。与Select不同，当向量中的某些通道具有不同条件时， if_then_else无法被向量化。

tvm.tir.if_then_else()的别名

tvm.te.isnan(x, span=None)

检查输入值是否为NaN。

Parameters:

x (PrimExpr) – 输入参数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

y – 结果。

Return type:

PrimExpr

tvm.tir.isnan()的别名

tvm.te.isfinite(x, span=None)

检查输入值是否为有限值。

Parameters:

x (PrimExpr) – 输入参数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

y – 结果。

Return type:

PrimExpr

tvm.tir.isfinite()的别名

tvm.te.isinf(x, span=None)

检查输入值是否为无穷大。

Parameters:

x (PrimExpr) – 输入参数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

y – 结果。

Return type:

PrimExpr

tvm.tir.isinf()的别名

tvm.te.div(a, b, span=None)

按照C/C++语义计算a / b。

Parameters:

a (PrimExpr) – 左侧操作数，已知为非负数。
b (PrimExpr) – 右侧操作数，已知为非负数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

res – 结果表达式。

Return type:

PrimExpr

注意

当操作数为整数时，返回 truncdiv(a, b, span)。

tvm.tir.div()的别名

tvm.te.indexdiv(a, b, span=None)

计算 floor(a / b)，其中 a 和 b 为非负数。

Parameters:

a (PrimExpr) – 左侧操作数，已知为非负数。
b (PrimExpr) – 右侧操作数，已知为非负数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

res – 结果表达式。

Return type:

PrimExpr

注意

使用此函数来分割非负索引。该函数可以利用操作数的非负性特性。

tvm.tir.indexdiv()的别名

tvm.te.indexmod(a, b, span=None)

计算indexdiv的余数。a和b为非负数。

Parameters:

a (PrimExpr) – 左侧操作数，已知为非负数。
b (PrimExpr) – 右侧操作数，已知为非负数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

res – 结果表达式。

Return type:

PrimExpr

注意

使用此函数来分割非负索引。该函数可以利用操作数的非负性特性。

tvm.tir.indexmod()的别名

tvm.te.truncdiv(a, b, span=None)

计算两个表达式的截断除法。

Parameters:

a (PrimExpr) – 左侧操作数
b (PrimExpr) – 右侧操作数
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

res – 结果表达式。

Return type:

PrimExpr

注意

这是C语言中默认的整数除法行为。

tvm.tir.truncdiv()的别名

tvm.te.truncmod(a, b, span=None)

计算两个表达式的截断模。

Parameters:

a (PrimExpr) – 左侧操作数
b (PrimExpr) – 右侧操作数
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

res – 结果表达式。

Return type:

PrimExpr

注意

这是C语言中默认的整数除法行为。

tvm.tir.truncmod()的别名

tvm.te.floordiv(a, b, span=None)

计算两个表达式的向下取整除。

Parameters:

a (PrimExpr) – 左侧操作数
b (PrimExpr) – 右侧操作数
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

res – 结果表达式。

Return type:

PrimExpr

tvm.tir.floordiv()的别名

tvm.te.floormod(a, b, span=None)

计算两个表达式的取模运算结果。

Parameters:

a (PrimExpr) – 左侧操作数
b (PrimExpr) – 右侧操作数
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

res – 结果表达式。

Return type:

PrimExpr

tvm.tir.floormod() 的别名

tvm.te.comm_reducer(fcombine, fidentity, name='reduce')

为归约操作创建一个可交换的归约器。

Parameters:

fcombine (function(Expr -> Expr -> Expr)) – 一个二元函数，接收两个Expr作为输入并返回一个Expr。
fidentity (function(str -> Expr)) – 一个接收类型字符串作为输入并返回常量表达式的函数。

Returns:

reducer – 一个用于创建轴上的归约表达式的函数。有两种使用方式：

接受 (expr, axis, where) 参数，在指定轴上生成 Reduce Expr；
简单地与多个 Exprs 一起使用。

Return type:

函数

示例

n = te.var("n")
m = te.var("m")
mysum = te.comm_reducer(lambda x, y: x+y,
    lambda t: tvm.tir.const(0, dtype=t), name="mysum")
A = te.placeholder((n, m), name="A")
k = te.reduce_axis((0, m), name="k")
B = te.compute((n,), lambda i: mysum(A[i, k], axis=k), name="B")

tvm.tir.comm_reducer()的别名

tvm.te.min(expr, axis, where=None, init=None, *args)

在轴上创建一个最小表达式。

Parameters:

expr (PrimExpr) – 源表达式。
axis (IterVar) – 归约操作的IterVar轴
where (可选, Expr) – 归约操作的过滤谓词。

Returns:

value – 结果值。

Return type:

PrimExpr

示例

m = te.var("m")
n = te.var("n")
A = te.placeholder((m, n), name="A")
k = te.reduce_axis((0, n), name="k")

# there are two way to use this min reducer:
# mode 1, accept (expr, axis, where) to produce an Reduce Expr
# tvm.min represents tvm.te.min or tvm.tir.min.
B = te.compute((m,), lambda i: tvm.min(A[i, k], axis=k), name="B")

# mode 2, simply use it with multiple Exprs:
min_res = tvm.min(m, n)

tvm.tir.min()的别名

tvm.te.max(expr, axis, where=None, init=None, *args)

创建一个沿轴的最大值表达式。

Parameters:

expr (PrimExpr) – 源表达式。
axis (IterVar) – 归约操作的迭代变量轴
where (可选参数, Expr) – 归约操作的过滤谓词。

Returns:

value – 结果值。

Return type:

PrimExpr

示例

m = te.var("m")
n = te.var("n")
A = te.placeholder((m, n), name="A")
k = te.reduce_axis((0, n), name="k")

# there are two way to use this max reducer:
# mode 1, accept (expr, axis, where) to produce an Reduce Expr
# tvm.max represents tvm.te.max or tvm.tir.max.
B = te.compute((m,), lambda i: tvm.max(A[i, k], axis=k), name="B")

# mode 2, simply use it with multiple Exprs:
max_res = tvm.max(m, n)

tvm.tir.max()的别名

tvm.te.sum(expr, axis, where=None, init=None, *args)

创建一个沿轴求和的表达式。

Parameters:

expr (PrimExpr) – 源表达式。
axis (IterVar) – 归约操作的迭代变量轴
where (可选参数, Expr) – 归约操作的过滤谓词。

Returns:

value – 结果值。

Return type:

PrimExpr

示例

m = te.var("m")
n = te.var("n")
A = te.placeholder((m, n), name="A")
k = te.reduce_axis((0, n), name="k")

# there are two way to use this sum reducer:
# mode 1, accept (expr, axis, where) to produce an Reduce Expr
# tvm.sum represents tvm.te.sum or tvm.tir.sum.
B = te.compute((m,), lambda i: tvm.sum(A[i, k], axis=k), name="B")

# mode 2, simply use it with multiple Exprs:
sum_res = tvm.sum(m, n)

tvm.tir.sum()的别名

tvm.te.add(lhs, rhs, span=None)

通用加法运算符。

Parameters:

lhs (object) – 左操作数。
rhs (object) – 右操作数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

op – 加法运算的结果表达式。

Return type:

tvm.Expr

tvm.tir.add()的别名

tvm.te.subtract(lhs, rhs, span=None)

通用减法运算符。

Parameters:

lhs (object) – 左操作数。
rhs (object) – 右操作数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

op – 减法运算的结果表达式。

Return type:

tvm.Expr

tvm.tir.subtract()的别名

tvm.te.multiply(lhs, rhs, span=None)

通用乘法运算符。

Parameters:

lhs (object) – 左操作数。
rhs (object) – 右操作数。
span (可选[Span]) – 该操作符在源代码中的位置。

Returns:

op – 乘法运算的结果表达式。

Return type:

tvm.Expr

tvm.tir.multiply()的别名

class tvm.te.TensorSlice(tensor, indices)

用于支持张量切片语法的辅助数据结构。

方法：

asobject()

将切片转换为对象。

属性：

dtype

张量的数据内容。

asobject(): 将切片转换为对象。

property dtype: 张量的数据内容。

class tvm.te.Tensor

张量对象，要构造它，请参见function.Tensor

属性：

`ndim`	张量的维度。
`axis`	张量的轴。
`op`	对应的 `Operation`。
`value_index`	输出值索引对应的张量。
`shape`	张量的输出形状。

property ndim: 张量的维度。

property axis: 张量的轴。

property op: 对应的 Operation。

property value_index: 张量对应的输出值索引。

property shape: 张量的输出形状。

tvm.te.tag_scope(tag)

操作符标签的作用范围。

Parameters:: 标签 (str) – 标签名称。
Returns:: tag_scope – 标签作用域对象，可用作装饰器或上下文管理器。
Return type:: TagScope

示例

n = te.var('n')
m = te.var('m')
l = te.var('l')
A = te.placeholder((n, l), name='A')
B = te.placeholder((m, l), name='B')
k = te.reduce_axis((0, l), name='k')

with tvm.te.tag_scope(tag='matmul'):
    C = te.compute((n, m), lambda i, j: te.sum(A[i, k] * B[j, k], axis=k))

# or use tag_scope as decorator
@tvm.te.tag_scope(tag="conv")
def compute_relu(data):
    return te.compute(data.shape, lambda *i: tvm.tir.Select(data(*i) < 0, 0.0, data(*i)))

tvm.te.placeholder(shape, dtype=None, name='placeholder')

构造一个空的张量对象。

Parameters:

shape (Tuple of Expr) – 张量的形状
dtype (str, optional) – 张量的数据类型
name (str, optional) – 张量的名称提示

Returns:

tensor – 创建的张量

Return type:

Tensor

tvm.te.compute(shape, fcompute, name='compute', tag='', attrs=None, varargs_names=None)

通过计算形状域构造一个新的张量。

计算规则是 result[axis] = fcompute(axis)

Parameters:

shape (Tuple of Expr) – 张量的形状
fcompute (lambda函数，形式为 indices-> value) – 指定输入源表达式
name (str, optional) – 张量的名称提示
tag (str, optional) – 关于计算的附加标签信息。
attrs (dict, optional) – 关于计算的额外辅助属性。
varargs_names (list, optional) – 用于每个可变参数的名称。如果未提供，可变参数将默认命名为i1、i2、...

Returns:

tensor – 创建的张量

Return type:

Tensor

tvm.te.scan(init, update, state_placeholder, inputs=None, name='scan', tag='', attrs=None)

通过沿轴扫描构建新的张量。

Parameters:

init (Tensor 或 list 的 Tensor) – 前init.shape[0]个时间戳的初始条件
update (Tensor 或 list 的 Tensor) – 由符号张量给出的扫描更新规则。
state_placeholder (Tensor 或 Tensor 列表) – 由update使用的占位变量。
inputs (Tensor 或 list 包含 Tensor，可选) – 扫描操作的输入列表。这不是必需的，但可以帮助编译器更快地检测扫描体。
name (str, optional) – 张量的名称提示
tag (str, optional) – 关于计算的附加标签信息。
attrs (dict, optional) – 关于计算的额外辅助属性。

Returns:

tensor – 创建的张量或张量元组包含多个输出。

Return type:

Tensor 或 list 的 Tensors

示例

# The following code is equivalent to numpy.cumsum
m = te.var("m")
n = te.var("n")
X = te.placeholder((m, n), name="X")
s_state = te.placeholder((m, n))
s_init = te.compute((1, n), lambda _, i: X[0, i])
s_update = te.compute((m, n), lambda t, i: s_state[t-1, i] + X[t, i])
res = tvm.te.scan(s_init, s_update, s_state, X)

tvm.te.extern(shape, inputs, fcompute, name='extern', dtype=None, in_buffers=None, out_buffers=None, tag='', attrs=None)

通过外部函数计算多个张量。

Parameters:

shape (元组或列表的元组) – 输出的形状。
inputs (list of Tensor) – 输入张量列表
fcompute (lambda函数，参数为inputs和outputs-> stmt) –
指定执行计算的IR语句。关于fcompute的函数签名，请参阅以下说明
注意

参数
- ins (tvm.tir.Buffer列表) - 每个输入的占位符
- outs (tvm.tir.Buffer列表) - 每个输出的占位符
返回值
- stmt (tvm.tir.Stmt) - 执行数组计算的语句。
name (str, optional) – 张量的名称提示
dtype (str 或 list 包含 str，可选) – 输出的数据类型，默认情况下dtype将与输入保持一致。
in_buffers (tvm.tir.Buffer 或 list 的 tvm.tir.Buffer, 可选) – 输入缓冲区。
out_buffers (tvm.tir.Buffer or list of tvm.tir.Buffer, optional) – 输出缓冲区。

tag: str, optional: 关于计算的附加标签信息。
attrs: dict, optional: 关于计算的额外辅助属性。

Returns:: tensor – 创建的张量或张量元组包含多个输出。
Return type:: Tensor 或 list 的 Tensors

示例

在下面的代码中，C是通过调用外部PackedFunc tvm.contrib.cblas.matmul生成的

A = te.placeholder((n, l), name="A")
B = te.placeholder((l, m), name="B")
C = te.extern((n, m), [A, B],
               lambda ins, outs: tvm.tir.call_packed(
                  "tvm.contrib.cblas.matmul",
                    ins[0], ins[1], outs[0], 0, 0), name="C")

tvm.te.var(name='tindex', dtype='int32', span=None)

创建一个具有指定名称和数据类型的新变量

Parameters:

name (str) – 名称
dtype (str) – 数据类型
span (可选[Span]) – 该变量在源代码中的位置。

Returns:

var – 结果符号变量。

Return type:

tir.Var

tvm.te.size_var(name='size', dtype='int32', span=None)

创建一个新变量表示张量形状大小，该大小必须为非负数。

Parameters:

name (str) – 名称
dtype (str) – 数据类型
span (可选[Span]) – 该变量在源代码中的位置。

Returns:

var – 结果符号形状变量。

Return type:

SizeVar

tvm.te.const(value, dtype='int32', span=None)

创建一个具有指定值和数据类型的新常量

Parameters:

value (Union[bool, int, float, numpy.ndarray, tvm.nd.NDArray]) – 常量值。
dtype (str) – 数据类型
span (可选[Span]) – 该变量在源代码中的位置。

Returns:

const – 结果常量表达式。

Return type:

PrimExpr

tvm.te.thread_axis(dom=None, tag='', name='', span=None)

创建一个新的IterVar来表示线程索引。

Parameters:

dom (Range 或 str) – 迭代的域当传入str时，dom会被设为None且str会被用作标签
tag (str, optional) – 线程标签
name (str, optional) – 变量的名称。
span (可选[Span]) – 该变量在源代码中的位置。

Returns:

axis – 线程迭代变量。

Return type:

IterVar

tvm.te.reduce_axis(dom, name='rv', thread_tag='', span=None)

创建一个新的归约迭代变量。

Parameters:

dom (Range) – 迭代的域。
name (str) – 变量的名称。
thread_tag (可选[str]) - 线程标签的名称。
span (可选[Span]) – 该变量在源代码中的位置。

Returns:

axis - 一个表示值的迭代变量。

Return type:

IterVar

tvm.te.create_prim_func(ops: 列表[张量 | 变量], index_dtype_override: str | None = None) → PrimFunc

从张量表达式创建TensorIR PrimFunc

Parameters:: ops (List[Union[_tensor.Tensor, tvm.tir.Var]]) – 源表达式。

示例

我们使用以下代码定义一个矩阵乘法内核：

import tvm
from tvm import te
from tvm.te import create_prim_func
import tvm.script

A = te.placeholder((128, 128), name="A")
B = te.placeholder((128, 128), name="B")
k = te.reduce_axis((0, 128), "k")
C = te.compute((128, 128), lambda x, y: te.sum(A[x, k] * B[y, k], axis=k), name="C")
func = create_prim_func([A, B, C])
print(func.script())

如果我们想使用TensorIR调度来对此类内核进行转换，需要通过create_prim_func([A, B, C])创建一个可调度的PrimFunc。生成的函数如下所示：

@T.prim_func
def tir_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (128, 128))
    B = T.match_buffer(b, (128, 128))
    C = T.match_buffer(c, (128, 128))

    for i, j, k in T.grid(128, 128, 128):
        with T.block():
            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
            with T.init():
                C[vi, vj] = 0.0
            C[vi, vj] += A[vi, vk] * B[vj, vk]

Returns:: func – 创建的函数。
Return type:: tir.PrimFunc

tvm.te.extern_primfunc(input_tensors: 列表[张量], primfunc: PrimFunc, **kwargs)

通过可调度的TIR PrimFunc计算张量

Parameters:

input_tensors (list of Tensor) - 输入张量，映射到对应的primfunc输入参数。
primfunc (PrimFunc) – TIR PrimFunc函数

Returns:

tensor – 创建的张量，如果包含多个输出则为张量元组。

Return type:

Tensor 或 list 的 Tensors

示例

在下面的代码中，一个TVMScript定义的TIR PrimFunc被内联到一个TE ExternOp中。对此应用te.create_prim_func

A = te.placeholder((128, 128), name="A")
B = te.placeholder((128, 128), name="B")

@T.prim_func
def before_split(a: T.handle, b: T.handle) -> None:
    A = T.match_buffer(a, (128, 128))
    B = T.match_buffer(b, (128, 128))
    for i, j in T.grid(128, 128):
        with T.block("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj] * 2.0

C = te.extern_primfunc([A, B], func)

class tvm.te.PlaceholderOp: 占位符操作。

class tvm.te.ComputeOp: 标量运算。

class tvm.te.ScanOp

扫描操作。

属性：

scan_axis

表示扫描轴，仅在作为ScanOp时定义

property scan_axis: 表示扫描轴，仅在作为ScanOp时定义

class tvm.te.ExternOp: 外部操作。