跳到主要内容
版本：0.21.0

tvm.te

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

函数:

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 取双曲 tanh。
sigmoid(x)快速获取 S 形函数。
log(x)对输入 x 取对数。
tan(x)对输入 x 取 tan。
cos(x)取输入 x 的 cos。
sin(x)对输入 x 取正弦值。
sqrt(x)对输入 x 取平方根。
rsqrt(x)取输入 x 的平方根的倒数。
floor(x[,span])取浮点输入 x 的下限。
ceil(x[,span])对浮点输入 x 取上限。
sinh(x)对输入 x 取 sinh。
cosh(x)对输入 x 取余弦值。
log2(x)对输入 x 取 log2。
log10(x)对输入 x 取 log10。
asin(x)取输入 x 的 asin。
asinh(x)取输入 x 的正弦值。
acos(x)对输入 x 取余数。
acosh(x)对输入 x 取余数。
atan(x)对输入 x 取正切值。
atanh(x)对输入 x 进行 atanh 处理。
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 的余数。a 和 b 非负。
truncdiv(a,b[,span])计算两个表达式的 truncdiv。
truncmod(a,b[,span])计算两个表达式的 truncmod。
floordiv(a,b[,span])计算两个表达式的 floordiv。
floormod(a,b[,span])计算两个表达式的 floormod。
logaddexp(a,b[,span])计算两个表达式的 logaddexp。
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[,name,thread_tag,span])创建一个新的 IterVar 进行缩减。
create_prim_func(ops[,index_dtype_override])从张量表达式创建 TensorIR PrimFunc。
extern_primfunc(input_tensors,primfunc,…)通过可调度的 TIR PrimFunc 计算张量。

类:

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

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

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

  • 参数：
    • args（列表）：符号布尔表达式列表。
    • span（可选[ Span]）：此运算符在源代码中的位置。
  • 返回：expr：表达式。
  • 返回类型： Expr。

别名tvm.tir.any()

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

创建一个新的表达式，该表达式表示所有参数条件的交集。

  • 参数：
    • args（列表）：符号布尔表达式列表。
    • span*（可选*[ Span]）：此运算符在源代码中的位置。
  • 返回：expr：表达式
  • 返回类型： Expr。

别名tvm.tir.all()

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

dtype 的最小值。

  • 参数：
    • dtype（str）：数据类型。
    • span*（可选*[ Span]）：此运算符在源代码中的位置。
  • 返回：value：dtype 的最小值。
  • 返回类型： tvm.Expr。

别名tvm.tir.min_value()

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

dtype 的最大值。

  • 参数：
    • dtype（str）*：*数据类型。
    • span*（可选*[ 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（Expr或 Buffers列表。）： 位置参数。
    • trace_action（str.）：跟踪操作的名称。
  • 返回：call：调用表达式。
  • 返回类型：PrimExpr。
另见

tvm.tir.call_packed

创建打包函数。

别名 tvm.tir.trace()

tvm.te.exp(x)​

取输入 x 的指数。

别名tvm.tir.exp()

tvm.te.erf(x)​

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

别名tvm.tir.erf()

tvm.te.tanh(x)​

对输入 x 取双曲 tanh。

别名tvm.tir.tanh()

tvm.te.sigmoid(x)​

快速获取 S 形函数。

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

别名tvm.tir.sigmoid()

tvm.te.log(x)​

对输入 x 取对数。

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

别名tvm.tir.log()

tvm.te.tan(x)​

对输入 x 取 tan。

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

别名tvm.tir.tan()

tvm.te.cos(x)​

取输入 x 的 cos。

  • 参数：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 的下限。

  • 参数：
    • x ( PrimExpr )：输入参数。
    • span*（可选*[ Span]）：此运算符在源代码中的位置。
  • 返回：y ：结果。
  • 返回类型：PrimExpr。

别名tvm.tir.floor()

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

对浮点输入 x 取上限。

  • 参数：
    • x ( PrimExpr )： 输入参数。
    • span*（可选*[ Span]）：此运算符在源代码中的位置。
  • 返回：y ：结果。
  • 返回类型：PrimExpr。

别名tvm.tir.ceil()

tvm.te.sinh(x)​

对输入 x 取 sinh。

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

别名tvm.tir.sinh()

tvm.te.cosh(x)​

对输入 x 取余弦值。

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

别名tvm.tir.cosh()

tvm.te.log2(x)​

对输入 x 取 log2。

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

别名tvm.tir.log2()

tvm.te.log10(x)​

对输入 x 取 log10。

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

别名tvm.tir.log10()

tvm.te.asin(x)​

取输入 x 的 asin。

  • 参数：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 进行 atanh 处理。

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

别名tvm.tir.atanh()

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

获取输入的截断值。

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

  • 参数：
    • x ( PrimExpr )：输入参数。
    • span*（可选*[ Span]）：此运算符在源代码中的位置。
  • 返回：y ：结果。
  • 返回类型：PrimExpr。

别名tvm.tir.trunc()

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

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

  • 参数：
    • x ( PrimExpr )：输入参数。
    • span*（可选*[ Span]）：此运算符在源代码中的位置。
  • 返回：y ：结果。
  • 返回类型：PrimExpr。

别名tvm.tir.abs()

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

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

  • 参数：
    • x ( PrimExpr )：输入参数。
    • span（可选*[ Span]）：此运算符在源代码中的位置。
  • 返回：y ：结果。
  • 返回类型：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。

  • 参数：
    • x ( PrimExpr )：输入参数。
    • span*（可选*[ Span]）：此运算符在源代码中的位置。
  • 返回：y ： 结果。
  • 返回类型：PrimExpr。

别名tvm.tir.nearbyint()

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

x 次方 y。

  • 参数：
    • x ( PrimExpr )：输入参数。
    • y ( PrimExpr )：指数。
    • span*（可选*[ 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 相同。

  • 参数：
  • 返回：z ：结果。
  • 返回类型：PrimExpr。

别名tvm.tir.fmod()

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

条件选择表达式。

  • 参数：
    • cond ( PrimExpr )：条件。
    • t ( PrimExpr )： 如果 cond 为真，则结果表达式。
    • f ( PrimExpr )：如果 cond 为假，则结果表达式。
    • span*（可选*[ Span]）：此运算符在源中的位置。
  • 返回：result：条件表达式的结果。
  • 返回类型：Node。

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

别名tvm.tir.if_then_else()

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

检查输入值是否为 Nan。

  • 参数：
    • x ( PrimExpr )： 输入参数。
    • span*（可选*[ Span]）：此运算符在源代码中的位置。
  • 返回：y ：结果。
  • 返回类型：PrimExpr。

别名tvm.tir.isnan()

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

检查输入值是否有限。

  • 参数：
    • x ( PrimExpr )：输入参数。
    • span*（可选*[ Span]）：此运算符在源代码中的位置。
  • 返回：y ： 结果。
  • 返回类型：PrimExpr。

别名tvm.tir.isfinite()

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

检查输入值是否无限。

  • 参数：
    • x ( PrimExpr )：输入参数。
    • span（可选[ Span]）：此运算符在源代码中的位置。
  • 返回：y ：结果。
  • 返回类型：PrimExpr。

别名tvm.tir.isinf()

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

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

  • 参数：
    • a ( PrimExpr )：左侧操作数，已知为非负数。
    • b ( PrimExpr )：右侧操作数，已知为非负。
    • span（可选[ 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*（可选*[ Span]）：此运算符在源中的位置。
  • 返回：res：结果表达式。
  • 返回类型：PrimExpr。

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

别名tvm.tir.indexdiv()

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

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

  • 参数：
    • a ( PrimExpr )：左侧操作数，已知为非负数。
    • b ( PrimExpr )：右侧操作数，已知为非负。
    • span*（可选*[ Span]）：此运算符在源中的位置。
  • 返回：res：结果表达式。
  • 返回类型：PrimExpr。

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

别名tvm.tir.indexmod()

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

计算两个表达式的 truncdiv。

  • 参数：
    • a ( PrimExpr )：左侧操作数。
    • b ( PrimExpr )：右侧操作数。
    • span*（可选*[ Span]）：此运算符在源中的位置。
  • 返回：res：结果表达式。
  • 返回类型：PrimExpr。

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

别名tvm.tir.truncdiv()

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

计算两个表达式的 truncmod。

  • 参数：
    • a ( PrimExpr )：左侧操作数。
    • b ( PrimExpr )：右侧操作数。
    • span（可选[Span]）：此运算符在源中的位置。
  • 返回：res：结果表达式。
  • 返回类型：PrimExpr。

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

别名tvm.tir.truncmod()

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

计算两个表达式的 floordiv。

  • 参数：
    • a ( PrimExpr )：左侧操作数。
    • b ( PrimExpr )：右侧操作数。
    • span*（可选*[ Span]）：此运算符在源中的位置。
  • 返回：res：结果表达式。
  • 返回类型：PrimExpr。

别名tvm.tir.floordiv()

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

计算两个表达式的 floormod。

  • 参数：
    • a ( PrimExpr )：左侧操作数。
    • b ( PrimExpr )：右侧操作数。
    • span（可选[ Span]）：此运算符在源中的位置。
  • 返回：res：结果表达式。
  • 返回类型：PrimExpr。

别名tvm.tir.floormod()

tvm.te.logaddexp(a, b, span=None)​

计算两个表达式的 logaddexp。

  • 参数：
    • a ( PrimExpr )：左侧操作数。
    • b ( PrimExpr )： 右侧操作数。
    • span（可选[Span]）：此运算符在源中的位置。
  • 返回：res：结果表达式。
  • 返回类型：PrimExpr。

别名tvm.tir.logaddexp()

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

创建一个交换减速器用于减速。

  • 参数：
    • fcombine (function ( Expr -> Expr -> Expr)) ：一个二元函数，以两个 Expr 作为输入并返回一个 Expr。
    • fidentity (function ( str -> Expr)) ：以字符串类型作为输入并返回 const Expr 的函数。
  • 返回：reducer：在 axis 上创建 Reduce 表达式的函数。有两种使用方法：
    • accept (expr, axis, where) 来在指定轴上生成一个 Reduce Expr；
    • 直接使用多个 Exprs。
  • 返回类型：函数

示例

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）：源表达式。
    • 轴（IterVar）：缩减 IterVar 轴
    • where（可选，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")

# 使用这个最小值归约器（min reducer）有两种方式：
# 方式 1：接受 (expr, axis, where) 参数来生成一个归约表达式（Reduce Expr）
# tvm.min 表示 tvm.te.min 或 tvm.tir.min。

B = te.compute((m,), lambda i: tvm.min(A[i, k], axis=k), name="B")

# 方式 2：直接用于多个表达式：
min_res = tvm.min(m, n)

别名tvm.tir.min()

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

在轴上创建最大表达式。

  • 参数：
    • expr（PrimExpr）：源表达式。
    • 轴（IterVar）：缩减 IterVar 轴。
    • where*（可选*，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")

# 使用这个最大值归约器（max reducer）有两种方式：
# 方式 1：接受 (expr, axis, where) 参数来生成一个归约表达式（Reduce Expr）
# tvm.max 表示 tvm.te.max 或 tvm.tir.max。

B = te.compute((m,), lambda i: tvm.max(A[i, k], axis=k), name="B")

# 方式 2：直接用于多个表达式：
max_res = tvm.max(m, n)

别名tvm.tir.max()

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

在轴上创建一个求和表达式。

  • 参数：
    • expr（PrimExpr）：源表达式。
    • 轴（IterVar）：缩减 IterVar 轴。
    • where*（可选*，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")

# 使用这个求和归约器（sum reducer）有两种方式：
# 方式 1：接受 (expr, axis, where) 参数来生成一个归约表达式（Reduce Expr）
# tvm.sum 表示 tvm.te.sum 或 tvm.tir.sum。

B = te.compute((m,), lambda i: tvm.sum(A[i, k], axis=k), name="B")

# 方式 2：直接用于多个表达式：
sum_res = tvm.sum(m, n)

别名tvm.tir.sum()

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

通用加法运算符。

  • 参数：
    • lhs（对象）：左操作数。
    • rhs（对象）：右操作数。
    • span*（可选*[ Span]）：此运算符在源中的位置。
  • 返回：op：加法运算的结果 Expr。
  • 返回类型： tvm.Expr。

别名tvm.tir.add()

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

通用减法运算符。

  • 参数：
    • lhs（对象）：左操作数。
    • rhs（对象）：右操作数。
    • span*（可选*[ Span]）：此运算符在源中的位置。
  • 返回：op：减法运算的结果 Expr。
  • 返回类型： tvm.Expr。

别名tvm.tir.subtract()

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

通用乘法运算符。

  • 参数：
    • lhs（对象）：左操作数。
    • rhs（对象）： 右操作数。
    • span*（可选*[ Span]）： 此运算符在源中的位置。
  • 返回：op：乘法运算的结果 Expr。
  • 返回类型： tvm.Expr。

别名tvm.tir.multiply()

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

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

方法：

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

属性：

dtype张量的数据内容。

asobject()​

将切片转换为对象。

property dtype​

张量的数据内容。

class tvm.te.Tensor​

Tensor 对象，构造方法参见 function.Tensor

属性：

ndim张量的维度。

property ndim​

张量的维度。

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

# 或者使用 tag_scope 作为装饰器 @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（Expr的元组）：张量的形状。
    • dtype（str，可选）：张量的数据类型。
    • name（str，可选）：张量的名称提示。
  • 返回：tensor：创建的张量。
  • 返回类型：Tensor。

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

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

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

  • 参数：
    • shape（Expr的元组）：张量的形状。
    • fcompute（indices–>value 的lambda 函数）：指定输入源表达式。
    • name（str，可选）： 张量的名称提示。
    • tag（str，可选）：有关计算的附加标签信息。
    • attrs（dict，可选）： 有关计算的附加辅助属性。
    • varargs_names（list，可选）：每个可变参数使用的名称。如果未提供，可变参数将被称为 i1、i2……
  • 返回：tensor：创建的张量。
  • 返回类型：Tensor。

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

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

  • 参数：
    • init ( Tensor或Tensor列表)：第一个 init.shape[0] 时间戳的初始条件。
    • 更新（Tensor或Tensor列表）：由符号张量给出的扫描更新规则。
    • state_placeholder（Tensor或Tensor列表）：更新使用的占位符变量。
    • 输入( Tensor 或 Tensor 列表，可选)：扫描的输入列表。这不是必需的，但有助于编译器更快地检测扫描主体。
    • name（str，可选）：张量的名称提示。
    • tag（str，可选）：有关计算的附加标签信息。
    • attrs（dict，可选）：有关计算的附加辅助属性。
  • 返回：tensor：创建的张量或张量元组包含多个输出。
  • 返回类型：Tensor or list of Tensors。

示例

# 以下代码等价于 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（元组或元组列表。）：输出的形状。
  • 输入（Tensor列表）：输入
  • fcompute*（输入的lambda 函数，输出–> stmt）： 指定用于执行计算的 IR 语句。请参阅以下注释以了解 fcompute 的函数签名。
  • name（str，可选）： 张量的名称提示。
  • dtype（str或str列表，可选）：输出的数据类型，默认情况下 dtype 与输入相同。
  • in_buffers（tvm.tir.Buffer或tvm.tir.Buffer列表，可选）：输入缓冲区。
  • out_buffers（tvm.tir.Buffer或tvm.tir.Buffer列表，可选）：输出缓冲区。

tag: str, optional

有关计算的附加标签信息。

attrs: dict, optional

有关计算的附加辅助属性。

  • 返回：tensor：创建的张量或张量元组包含多个输出。
  • 返回类型：Tensor or 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)​

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

  • 参数：
    • name（str）：名称。
    • dtype ( str )：数据类型。
    • span*（可选*[ Span]）：此变量在源中的位置。
  • 返回：var：结果符号变量。
  • 返回类型：tir.Var。

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

创建一个新变量表示张量形状的大小，它是非负的。

  • 参数：
    • name（str）：名称。
    • dtype ( str )： 数据类型。
    • span*（可选*[Span]）：此变量在源中的位置。
  • 返回：var：结果符号形状变量。
  • 返回类型：SizeVar。

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

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

  • 参数：
    • value ( Union[bool,int,float,numpy.ndarray,tvm.nd.NDArray] )： 常量值。
    • dtype ( str )：数据类型。
    • span*（可选*[ Span]）： 此变量在源中的位置。
  • 返回：const：结果常量 expr。
  • 返回类型：PrimExpr。

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

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

  • 参数：
    • dom ( Range或str )： 迭代的范围，当传入 str 时，dom 设置为 None，str 作为标签。
    • tag（str，可选）：线程标签。
    • name（str，可选）：变量的名称。
    • span（可选[Span]）： 此变量在源中的位置。
  • 返回：axis：线程迭代变量。
  • 返回类型：IterVar。

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

创建一个新的 IterVar 进行缩减。

  • 参数：
    • dom (范围)： 迭代的范围。
    • name ( str )： 变量的名称。
    • thread_tag*（可选*[ str]）：thread_tag 的名称。
    • span*（可选*[ 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 核：

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]
  • 返回：func：创建的函数。
  • 返回类型：tir.PrimFunc

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

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

  • 参数：
    • input_tensors（Tensor列表）：映射到相应 primfunc 输入参数的输入张量。
    • primfunc ( PrimFunc )：TIR PrimFunc。
  • 返回：tensor：如果包含多个输出，则创建的张量或张量元组。
  • 返回类型：Tensor or 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​

外部操作。

编辑此页
最后由 wenjing-user 于 更新