tvm.te

Tensor 表达式语言的命名空间

函数

`any`(*args[, span])	创建由参数中所有条件的并集组成的新表达式
`all`(*args[, span])	创建由参数中所有条件的交集组成的新表达式
`min_value`(dtype[, span])	dtype 的最小值
`max_value`(dtype[, span])	dtype 的最大值
`trace`(args[, trace_action])	在运行时跟踪张量数据。
`exp`(x)	计算输入 x 的指数。
`erf`(x)	计算输入 x 的高斯误差函数。
`tanh`(x)	计算输入 x 的双曲正切。
`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 的 floor 值。
`ceil`(x[, span])	计算浮点输入 x 的 ceil 值。
`sinh`(x)	计算输入 x 的双曲正弦。
`cosh`(x)	计算输入 x 的双曲余弦。
`log2`(x)	计算输入 x 的以 2 为底的对数。
`log10`(x)	计算输入 x 的以 10 为底的对数。
`asin`(x)	计算输入 x 的反正弦。
`asinh`(x)	计算输入 x 的反双曲正弦。
`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 中设置的位数。
`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])	计算两个表达式的 truncdiv。
`truncmod`(a, b[, span])	计算两个表达式的 truncmod。
`floordiv`(a, b[, span])	计算两个表达式的 floordiv。
`floormod`(a, b[, span])	计算两个表达式的 floormod。
`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])	创建具有指定名称和 dtype 的新变量
`size_var`([name, dtype, span])	创建一个表示张量形状大小的新变量，该变量为非负数。
`const`(value[, dtype, span])	创建具有指定值和 dtype 的新常量
`thread_axis`([dom, tag, name, span])	创建一个新的 IterVar 来表示线程索引。
`reduce_axis`(dom[, name, thread_tag, span])	创建一个用于规约的新 IterVar。
`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)

创建由参数中所有条件的并集组成的新表达式

参数:

args (list) – 符号布尔表达式列表
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

expr – 表达式

返回类型:

Expr

别名：tvm.tir.any()

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

创建由参数中所有条件的交集组成的新表达式: 参数

参数:

args (list) – 符号布尔表达式列表
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

expr – 表达式

返回类型:

Expr

别名：tvm.tir.all()

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

dtype 的最小值

参数:

dtype (str) – 数据类型。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

value – dtype 的最小值。

返回类型:

tvm.Expr

别名：tvm.tir.min_value()

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

dtype 的最大值

参数:

dtype (str) – 数据类型。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

value – dtype 的最大值。

返回类型:

tvm.Expr

别名：tvm.tir.max_value()

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

在运行时跟踪张量数据。

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

参数:

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

返回:

call – 调用表达式。

返回类型:

PrimExpr

另请参阅

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

别名：tvm.tir.trace()

tvm.te.exp(x)

计算输入 x 的指数。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.exp()

tvm.te.erf(x)

计算输入 x 的高斯误差函数。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.erf()

tvm.te.tanh(x)

计算输入 x 的双曲正切。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.tanh()

tvm.te.sigmoid(x)

获取 sigmoid 的快速函数

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.sigmoid()

tvm.te.log(x)

计算输入 x 的对数。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.log()

tvm.te.tan(x)

计算输入 x 的正切。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.tan()

tvm.te.cos(x)

计算输入 x 的余弦。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.cos()

tvm.te.sin(x)

计算输入 x 的正弦。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.sin()

tvm.te.sqrt(x)

计算输入 x 的平方根。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.sqrt()

tvm.te.rsqrt(x)

计算输入 x 平方根的倒数。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.rsqrt()

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

计算浮点输入 x 的 floor 值。

参数:

x (PrimExpr) – 输入参数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

y – 结果。

返回类型:

PrimExpr

别名：tvm.tir.floor()

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

计算浮点输入 x 的 ceil 值。

参数:

x (PrimExpr) – 输入参数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

y – 结果。

返回类型:

PrimExpr

别名：tvm.tir.ceil()

tvm.te.sinh(x)

计算输入 x 的双曲正弦。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.sinh()

tvm.te.cosh(x)

计算输入 x 的双曲余弦。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.cosh()

tvm.te.log2(x)

计算输入 x 的以 2 为底的对数。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.log2()

tvm.te.log10(x)

计算输入 x 的以 10 为底的对数。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.log10()

tvm.te.asin(x)

计算输入 x 的反正弦。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.asin()

tvm.te.asinh(x)

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

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.asinh()

tvm.te.acos(x)

计算输入 x 的反余弦。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.acos()

tvm.te.acosh(x)

计算输入 x 的反余弦。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.acosh()

tvm.te.atan(x)

计算输入 x 的反正切。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.atan()

tvm.te.atanh(x)

计算输入 x 的反双曲正切。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

别名：tvm.tir.atanh()

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

获取输入的截断值。

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

参数:

x (PrimExpr) – 输入参数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

y – 结果。

返回类型:

PrimExpr

别名：tvm.tir.trunc()

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

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

参数:

x (PrimExpr) – 输入参数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

y – 结果。

返回类型:

PrimExpr

别名：tvm.tir.abs()

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

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

参数:

x (PrimExpr) – 输入参数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

y – 结果。

返回类型:

PrimExpr

别名：tvm.tir.round()

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

将数组元素四舍五入到最接近的整数。此内部函数使用 llvm.nearbyint 而不是 llvm.round，前者速度更快，但结果与 te.round 不同。值得注意的是，nearbyint 根据舍入模式进行舍入，而 te.round (llvm.round) 则忽略舍入模式。有关两者之间差异的详细信息，请参阅： https://cppreference.cn/w/cpp/numeric/math/round https://cppreference.cn/w/cpp/numeric/math/nearbyint

参数:

x (PrimExpr) – 输入参数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

y – 结果。

返回类型:

PrimExpr

tvm.tir.nearbyint() 的别名

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

x 的 y 次方

参数:

x (PrimExpr) – 输入参数。
y (PrimExpr) – 指数
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

z – 结果。

返回类型:

PrimExpr

tvm.tir.power() 的别名

tvm.te.popcount(x)

计算输入 x 中设置的位数。

参数:: x (PrimExpr) – 输入参数。
返回:: y – 结果。
返回类型:: PrimExpr

tvm.tir.popcount() 的别名

tvm.te.fmod(x, y)

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

参数:

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

返回:

z – 结果。

返回类型:

PrimExpr

tvm.tir.fmod() 的别名

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

条件选择表达式。

参数:

cond (PrimExpr) – 条件
t (PrimExpr) – 如果 cond 为 true，则为结果表达式。
f (PrimExpr) – 如果 cond 为 false，则为结果表达式。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

result – 条件表达式的结果。

返回类型:

节点

注意

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

tvm.tir.if_then_else() 的别名

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

检查输入值是否为 Nan。

参数:

x (PrimExpr) – 输入参数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

y – 结果。

返回类型:

PrimExpr

tvm.tir.isnan() 的别名

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

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

参数:

x (PrimExpr) – 输入参数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

y – 结果。

返回类型:

PrimExpr

tvm.tir.isfinite() 的别名

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

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

参数:

x (PrimExpr) – 输入参数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

y – 结果。

返回类型:

PrimExpr

tvm.tir.isinf() 的别名

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

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

参数:

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

返回:

res – 结果表达式。

返回类型:

PrimExpr

注意

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

tvm.tir.div() 的别名

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

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

参数:

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

返回:

res – 结果表达式。

返回类型:

PrimExpr

注意

使用此函数来拆分非负索引。此函数可能会利用操作数的非负性。

tvm.tir.indexdiv() 的别名

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

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

参数:

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

返回:

res – 结果表达式。

返回类型:

PrimExpr

注意

使用此函数来拆分非负索引。此函数可能会利用操作数的非负性。

tvm.tir.indexmod() 的别名

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

计算两个表达式的 truncdiv。

参数:

a (PrimExpr) – 左侧操作数
b (PrimExpr) – 右侧操作数
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

res – 结果表达式。

返回类型:

PrimExpr

注意

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

tvm.tir.truncdiv() 的别名

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

计算两个表达式的 truncmod。

参数:

a (PrimExpr) – 左侧操作数
b (PrimExpr) – 右侧操作数
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

res – 结果表达式。

返回类型:

PrimExpr

注意

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

tvm.tir.truncmod() 的别名

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

计算两个表达式的 floordiv。

参数:

a (PrimExpr) – 左侧操作数
b (PrimExpr) – 右侧操作数
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

res – 结果表达式。

返回类型:

PrimExpr

tvm.tir.floordiv() 的别名

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

计算两个表达式的 floormod。

参数:

a (PrimExpr) – 左侧操作数
b (PrimExpr) – 右侧操作数
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

res – 结果表达式。

返回类型:

PrimExpr

tvm.tir.floormod() 的别名

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

为规约创建可交换的规约器。

参数:

fcombine (function(Expr -> Expr -> Expr)) – 一个二元函数，它接受两个 Expr 作为输入以返回一个 Expr。
fidentity (function(str -> Expr)) – 一个函数，它接受一个类型字符串作为输入以返回一个常量 Expr。

返回:

reducer – 一个函数，它在轴上创建归约表达式。有两种使用方法

接受 (expr, axis, where) 以在指定的轴上生成 Reduce Expr；
简单地将其与多个 Expr 一起使用。

返回类型:

函数

示例

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)

创建轴上的最小值表达式。

参数:

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

返回:

value – 结果值。

返回类型:

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)

创建轴上的最大值表达式。

参数:

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

返回:

value – 结果值。

返回类型:

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)

创建轴上的总和表达式。

参数:

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

返回:

value – 结果值。

返回类型:

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)

通用加法运算符。

参数:

lhs (object) – 左操作数。
rhs (object) – 右操作数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

op – 加法运算的结果 Expr。

返回类型:

tvm.Expr

tvm.tir.add() 的别名

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

通用减法运算符。

参数:

lhs (object) – 左操作数。
rhs (object) – 右操作数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

op – 减法运算的结果 Expr。

返回类型:

tvm.Expr

tvm.tir.subtract() 的别名

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

通用乘法运算符。

参数:

lhs (object) – 左操作数。
rhs (object) – 右操作数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

op – 乘法运算的结果 Expr。

返回类型:

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)

运算符标签作用域。

参数:: tag (str) – 标签名称。
返回:: tag_scope – 标签作用域对象，可以用作装饰器或上下文管理器。
返回类型:: 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')

构造一个空的张量对象。

参数:

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

返回:

tensor – 创建的张量

返回类型:

Tensor

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

通过在形状域上计算来构造新张量。

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

参数:

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

返回:

tensor – 创建的张量

返回类型:

Tensor

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

通过沿轴扫描来构造新张量。

参数:

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

返回:

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

返回类型:

Tensor 或 list of 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)

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

参数:

shape (tuple 或 list of tuples.) – 输出的形状。
inputs (list of Tensor) – 输入
fcompute (lambda function of inputs, outputs-> stmt) –
指定执行计算的 IR 语句。有关 fcompute 的函数签名，请参见以下注释
注意

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

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

返回:: tensor – 创建的张量或包含多个输出的张量元组。
返回类型:: Tensor 或 list of 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)

创建具有指定名称和 dtype 的新变量

参数:

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

返回:

var – 结果符号变量。

返回类型:

tir.Var

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

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

参数:

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

返回:

var – 结果符号形状变量。

返回类型:

SizeVar

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

创建具有指定值和 dtype 的新常量

参数:

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

返回:

const – 结果常量表达式。

返回类型:

PrimExpr

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

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

参数:

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

返回:

axis – 线程迭代变量。

返回类型:

IterVar

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

创建一个用于规约的新 IterVar。

参数:

dom (Range) – 迭代的域。
name (str) – 变量的名称。
thread_tag (Optional[str]) – thread_tag 的名称。
span (Optional[Span]) – 此变量在源代码中的位置。

返回:

axis – 表示值的迭代变量。

返回类型:

IterVar

tvm.te.create_prim_func(ops: List[Tensor | Var], index_dtype_override: str | None = None) → PrimFunc

从张量表达式创建 TensorIR PrimFunc

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

示例

我们使用以下代码定义一个 matmul kernel

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 schedule 对此类 kernel 进行转换，我们需要使用 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]

返回:: func – 创建的函数。
返回类型:: tir.PrimFunc

tvm.te.extern_primfunc(input_tensors: List[Tensor], primfunc: PrimFunc, **kwargs)

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

参数:

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

返回:

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

返回类型:

Tensor 或 list of 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: 外部操作。