tvm.tir

张量级 IR 的命名空间

class tvm.tir.Buffer

TVM 中的符号数据缓冲区。

Buffer 提供了一种在 TVM 中表示数据结构的数据布局特例化的方法。

不要直接构造，请使用 decl_buffer() 代替。有关更多详细信息，请参阅 decl_buffer() 的文档。

另请参阅

decl_buffer: 声明缓冲区

access_ptr(access_mask, ptr_type='handle', content_lanes=1, offset=0, extent=None)

获取缓冲区头部的访问指针。

这是与外部函数交互时获取缓冲区数据指针的推荐方法。

参数:

access_mask (int) – 访问模式 MASK。指示访问将读取还是写入数据内容。
ptr_type (str, 可选) – 结果指针的数据类型。除非我们想将指针强制转换为特定类型，否则不要指定。
content_lanes (int, 可选) – 数据类型的通道数。对于向量类型，此值大于 1。
offset (Expr, 可选) – 指针的偏移量。我们可以使用它从 ptr 的地址偏移元素数量。
extent (Expr, 可选) – 指针的范围。

示例

# Get access ptr for read
buffer.access_ptr("r")
# Get access ptr for read/write with bitmask
buffer.access_ptr(Buffer.READ | Buffer.WRITE)
# Get access ptr for read/write with str flag
buffer.access_ptr("rw")
# Get access ptr for read with offset
buffer.access_ptr("r", offset = 100)
# Get access ptr for read with extent
buffer.access_ptr("r", extent = 100)

vload(begin, dtype=None, predicate=None)

生成一个从起始索引加载 dtype 的 Expr。

参数:

begin (Array of Expr) – 以 Buffer.dtype 为单位的起始索引
dtype (str) – 要加载的数据类型，可以是向量类型，其通道数是 Buffer.dtype 的倍数
predicate (Optional[PrimExpr]) – 一个布尔值向量掩码，指示要加载的向量的哪些通道。掩码的通道数必须等于要加载的通道数。

返回值:

load – 相应的加载表达式。

返回类型:

Expr

vstore(begin, value, predicate=None)

生成一个将值存储到起始索引的 Stmt。

参数:

begin (Array of Expr) – 以 Buffer.dtype 为单位的起始索引
value (Expr) – 要存储的值。
predicate (Optional[PrimExpr]) – 一个布尔值向量掩码，指示要存储的向量的哪些通道。掩码的通道数必须等于值中的通道数。

返回值:

store – 相应的存储 stmt。

返回类型:

Stmt

scope(): 返回与此缓冲区关联的存储范围。 :returns: scope – 与此缓冲区关联的存储范围。 :rtype: str

get_flattened_buffer()

生成一个作为此缓冲区扁平化版本的 Buffer。

返回值:: flattened – 相应的扁平缓冲区。
返回类型:: Buffer

offset_of(indices)

确定扁平缓冲区中提供的索引的偏移量。

参数:: indices (Union[PrimExpr, List[PrimExpr]]) – 原始缓冲区中元素的索引。
返回值:: flattened_indices – 扁平缓冲区中元素的偏移索引。
返回类型:: List[PrimExpr]

tvm.tir.decl_buffer(shape, dtype=None, name='buffer', data=None, strides=None, elem_offset=None, scope='', data_alignment=-1, offset_factor=0, buffer_type='', axis_separators=None, span=None)

声明一个新的符号缓冲区。

通常，缓冲区在 lower 和 build 期间自动创建。仅当用户想要指定自己的缓冲区布局时才需要这样做。

有关缓冲区用法的详细讨论，请参见下面的注释。

参数:

shape (tuple of Expr) – 缓冲区的形状。
dtype (str, 可选) – 缓冲区的数据类型。
name (str, 可选) – 缓冲区的名称。
data (tir.Var, 可选) – 缓冲区中的数据指针。
strides (array of Expr) – 缓冲区的步幅。
elem_offset (Expr, 可选) – 数组到数据的起始偏移量。以 dtype 的元素数量表示。
scope (str, 可选) – 缓冲区的存储范围，如果不是全局范围。如果 scope 等于空字符串，则表示它是全局内存。
data_alignment (int, 可选) – 数据指针的对齐（以字节为单位）。如果传递 -1，则对齐将设置为 TVM 的内部默认值。
offset_factor (int, 可选) – elem_offset 字段的因子，设置后，elem_offset 必须是 offset_factor 的倍数。如果传递 0，则对齐将设置为 1。如果传递非零值，如果 elem_offset 不为 None，我们将为 elem_offset 创建一个 tir.Var。
buffer_type (str, 可选, {"", "auto_broadcast"}) – auto_broadcast 缓冲区允许用户实现广播计算，而无需考虑维度大小是否等于 1。如果维度 j 的形状等于 1，则 TVM 将 buffer[i][j][k] 映射到 buffer[i][0][k]。
axis_separators (list of int, 可选) – 如果传递，则为轴组之间分隔符的列表，每个轴组都扁平化为输出轴。对于平面内存空间，应为 None 或空列表。
span (Optional[Span]) – decl_buffer 创建在源代码中的位置。

返回值:

buffer – 创建的缓冲区

返回类型:

tvm.tir.Buffer

注意

缓冲区数据结构反映了 dlpack 中的 DLTensor 结构。虽然 DLTensor 数据结构非常通用，但通常有助于创建仅处理数据结构特定情况的函数，并使编译后的函数从中受益。

如果用户在构造函数时将 strides 和 elem_offset 作为 None 传递，则该函数将专门用于紧凑且对齐的 DLTensor。如果用户将完全通用的符号数组传递给 strides，则生成的函数将变为完全通用。

class tvm.tir.DataProducer

class tvm.tir.Layout

布局由大写字母、小写字母和数字组成，其中大写字母表示主轴，相应的小写字母和因子大小表示从属轴。例如，NCHW16c 可以描述一个 5-D 张量 [batch_size, channel, height, width, channel_block]。这里从属轴 channel_block=16 是主轴 C（通道）的因子大小。

另请参阅

layout: 声明布局

index_of(axis)

获取轴的索引

参数:: axis (str) – 轴名称，需要为 [a-z,A-Z]
返回值:: index – 轴的索引，如果未找到则为 -1。
返回类型:: int

factor_of(axis)

获取从属轴的因子大小。

参数:: axis (str) – 轴名称，需要为 [a-z,A-Z]
返回值:: factor – 轴的从属轴的大小（如果 axis 是主轴），或轴本身的大小（如果 axis 是从属轴）。如果轴不在布局中，则返回 -1。
返回类型:: int

class tvm.tir.BijectiveLayout

两个布局（src-layout 和 dst-layout）的双射映射。它提供彼此之间的形状和索引转换。

不要直接构造，请使用 bijective_layout 代替。有关更多详细信息，请参阅 bijective_layout 的文档。

参数:

src_layout (str or Layout) – 源布局。
dst_layout (str or Layout) – 目标布局。

另请参阅

bijective_layout: 声明布局

forward_index(index)

给定 src-layout 的索引，推断 dst 索引。

参数:: index (Array of Expr) – src-layout 中的索引。
返回值:: dst_index – dst-layout 中推断的索引。
返回类型:: Array of Expr

backward_index(index)

给定 dst-layout 的索引，推断 src 索引。

参数:: index (Array of Expr) – dst-layout 中的索引。
返回值:: src_index – src-layout 中推断的索引。
返回类型:: Array of Expr

forward_shape(shape)

给定 src-layout 的形状，推断 dst 形状。

参数:: shape (Array of Expr) – src-layout 中的形状。
返回值:: dst_shape – dst-layout 中推断的形状。
返回类型:: Array of Expr

backward_shape(shape)

给定 dst-layout 的形状，推断 src 形状。

参数:: shape (Array of Expr) – dst-layout 中的形状。
返回值:: src_shape – src-layout 中推断的形状。
返回类型:: Array of Expr

tvm.tir.bijective_layout(src_layout: str | Layout, dst_layout: str | Layout) → BijectiveLayout

创建双射布局映射。

参数:

src_layout (str or Layout) – 源布局。
dst_layout (str or Layout) – 目标布局。

返回值:

bijective_layout – 创建的双射布局

返回类型:

BijectiveLayout

tvm.tir.layout(layout_str: str, dtype: str = 'int32') → Layout

从字符串创建布局节点。

参数:

layout_str (str) – 布局表示形式由大写字母、小写字母和数字组成，其中大写字母表示主轴，相应的小写字母和因子大小表示从属轴。例如，NCHW16c 可以描述一个 5-D 张量 [batch_size, channel, height, width, channel_block]。这里从属轴 channel_block=16 是主轴 C（通道）的因子大小。
dtype (str) – 返回的布局中生成的轴变量的 dtype。它必须是整数类型。

返回值:

layout – 创建的布局

返回类型:

Layout

class tvm.tir.Var(name: str, dtype: str | Type, span: Span | None = None)

符号变量。

参数:

name (str) – 名称
dtype (Union[str, ir.Type]) – 数据类型
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.SizeVar(name: str, dtype: str | Type, span: Span | None = None)

用于表示张量索引大小的符号变量: 它大于或等于零。

参数:

name (str) – 名称
dtype (Union[str, ir.Type]) – 数据类型
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Reduce(combiner: CommReducer, src: List[PrimExpr], rdom: List[IterVar], condition: PrimExpr, value_index: int, init: List[PrimExpr] | None = None, span: Span | None = None)

Reduce 节点。

参数:

combiner (CommReducer) – 合并器。
src (list of Expr) – 源表达式。
rdom (list of IterVar) – 迭代域
condition (PrimExpr) – reduce 条件。
value_index (int) – 值索引。
init (list of Expr) – 输出的初始值。这可以是 int、float 或 ProducerLoad
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.FloatImm(dtype: str, value: float, span: Span | None = None)

浮点常量。

参数:

dtype (str) – 数据类型
value (float) – 常数值。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.IntImm(dtype: str, value: int, span: Span | None = None)

整数常量。

参数:

dtype (str) – 数据类型
value (int) – 常量值。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.StringImm(value: str, span: Span | None = None)

字符串常量。

参数:

value (str) – 函数的值。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Cast(dtype, value, span: Span | None = None)

类型转换表达式。

参数:

dtype (str) – 数据类型
value (PrimExpr) – 函数的值。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Add(a: PrimExpr, b: PrimExpr, span: Span | None = None)

加法节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Sub(a: PrimExpr, b: PrimExpr, span: Span | None = None)

减法节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Mul(a: PrimExpr, b: PrimExpr, span: Span | None = None)

乘法节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Div(a: PrimExpr, b: PrimExpr, span: Span | None = None)

除法节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Mod(a: PrimExpr, b: PrimExpr, span: Span | None = None)

取模节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.FloorDiv(a: PrimExpr, b: PrimExpr, span: Span | None = None)

向下取整除法节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.FloorMod(a: PrimExpr, b: PrimExpr, span: Span | None = None)

向下取整取模节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Min(a: PrimExpr, b: PrimExpr, span: Span | None = None)

最小值节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Max(a: PrimExpr, b: PrimExpr, span: Span | None = None)

最大值节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.EQ(a: PrimExpr, b: PrimExpr, span: Span | None = None)

等于节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.NE(a: PrimExpr, b: PrimExpr, span: Span | None = None)

不等于节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.LT(a: PrimExpr, b: PrimExpr, span: Span | None = None)

小于节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.LE(a: PrimExpr, b: PrimExpr, span: Span | None = None)

小于等于节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.GT(a: PrimExpr, b: PrimExpr, span: Span | None = None)

大于节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.GE(a: PrimExpr, b: PrimExpr, span: Span | None = None)

大于等于节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.And(a: PrimExpr, b: PrimExpr, span: Span | None = None)

逻辑与节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Or(a: PrimExpr, b: PrimExpr, span: Span | None = None)

逻辑或节点。

参数:

a (PrimExpr) – 左操作数。
b (PrimExpr) – 右操作数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Not(a: PrimExpr, span: Span | None = None)

逻辑非节点。

参数:

a (PrimExpr) – 输入值
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Select(condition: PrimExpr, true_value: PrimExpr, false_value: PrimExpr, span: Span | None = None)

选择节点。

注意

Select 可能会计算 true_value 和 false_value。如果想要获取仅评估正确分支的条件表达式，请使用 tvm.tir.if_then_else 代替。

参数:

condition (PrimExpr) – 条件表达式。
true_value (PrimExpr) – 当条件为真时取的值。
false_value (PrimExpr) – 当条件为假时取的值。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.BufferLoad(buffer: Buffer, indices: List[PrimExpr], predicate: PrimExpr | None = None, span: Span | None = None)

Buffer 加载节点。

参数:

buffer (Buffer) – 要加载的 buffer。
indices (List[PrimExpr]) – 用于加载值的 buffer 索引。
span (Optional[Span]) – 此表达式在源代码中的位置。
predicate (Optional[PrimExpr]) – 一个布尔值向量掩码，指示要加载的向量的哪些通道。掩码的通道数必须等于要加载的通道数。

class tvm.tir.ProducerLoad(producer: DataProducer, indices: List[PrimExpr], span: Span | None = None)

Producer 加载节点。

参数:

producer (DataProducer) – 要加载的 buffer。
indices (List[PrimExpr]) – buffer 索引。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Ramp(base: PrimExpr, stride: PrimExpr, lanes: PrimExpr, span: Span | None = None)

Ramp 节点。

参数:

base (PrimExpr) – 基表达式。
stride (PrimExpr) – ramp 的步长。
lanes (PrimExpr) – 表达式的通道数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Broadcast(value: PrimExpr, lanes: PrimExpr, span: Span | None = None)

Broadcast 节点。

参数:

value (PrimExpr) – 表达式的值。
lanes (PrimExpr) – 表达式的通道数。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Shuffle(vectors: List[PrimExpr], indices: List[PrimExpr], span: Span | None = None)

Shuffle 节点。

参数:

vectors (List[PrimExpr]) – 向量
indices (List[PrimExpr]) – 索引
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Call(dtype: str, op: Op | str, args: List[PrimExpr], span: Span | None = None)

tir.Call 节点。

参数:

dtype (str) – 返回数据类型
op (Union[Op, str]) – 要调用的函数，或全局 tvm.Op 的名称
args (list of Expr) – 调用的输入参数
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.CallEffectKind: tir.Call 效果的可能种类。

class tvm.tir.Let(var: Var, value: PrimExpr, body: PrimExpr, span: Span | None = None)

Let 节点。

参数:

var (tir.Var) – 绑定中的变量。
value (PrimExpr) – 要绑定的值。
body (PrimExpr) – 主体表达式。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.IterVar(dom: Range, var: Var | str, iter_type: int, thread_tag: str = '', span: Span | None = None)

表示迭代变量。

IterVar 表示计算中的轴迭代。

参数:

dom (Range) – 迭代的域。
var (Union[tir.Var, str]) – 用于迭代的内部变量。
iter_type (int) – 迭代类型。
thread_tag (str) – 线程类型标签。
span (Optional[Span]) – 此表达式在源代码中的位置。

另请参阅

te.thread_axis: 创建线程轴 IterVar。
te.reduce_axis: 创建规约轴 IterVar。

class tvm.tir.CommReducer(lhs: List[Var], rhs: List[Var], result: List[PrimExpr], identity_element: List[PrimExpr], span: Span | None = None)

交换规约运算符

参数:

lhs (List[tir.Var]) – 规约器的左侧参数。
rhs (List[tir.Var]) – 规约器的右侧参数。
result (List[PrimExpr]) – 规约结果。
identity_element (List[PrimExpr]) – 单位元素。
span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.tir.Any(span: Span | None = None)

Any 节点。

spanOptional[Span]: 此表达式在源代码中的位置。

class tvm.tir.Stmt: 所有语句的基类。

class tvm.tir.LetStmt(var: Var, value: PrimExpr, body: Stmt, span: Span | None = None)

LetStmt 节点。

参数:

var (tir.Var) – 绑定中的变量。
value (PrimExpr) – 要绑定的值。
body (Stmt) – 主体语句。
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.AssertStmt(condition: PrimExpr, message: PrimExpr, body: Stmt, span: Span | None = None)

AssertStmt 节点。

参数:

condition (PrimExpr) – 断言条件。
message (PrimExpr) – 错误消息。
body (tvm.tir.Stmt) – 主体语句。
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.ForKind(value): for 循环的种类。

注意

ForKind 可以更改循环的控制流语义，需要在所有 TIR 通道中考虑。

class tvm.tir.For(loop_var: Var, min: PrimExpr, extent: PrimExpr, kind: ForKind, body: Stmt, thread_binding: IterVar | None = None, annotations: Mapping[str, Object] | None = None, span: Span | None = None)

For 节点。

参数:

loop_var (tir.Var) – 循环变量。
min (PrimExpr) – 起始值。
extent (PrimExpr) – 循环的长度。
kind (ForKind) – for 循环的类型。
body (Stmt) – 主体语句。
thread_binding (Optional[tir.IterVar]) – 此循环绑定到的线程。仅当 kind 为 ThreadBinding 时有效
annotations (Optional[Mapping[str, Object]]) – 附加的注解提示。
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.While(condition: PrimExpr, body: Stmt, span: Span | None = None)

While 节点。

参数:

condition (PrimExpr) – 终止条件。
body (Stmt) – 主体语句。
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.BufferStore(buffer: Buffer, value: PrimExpr, indices: List[PrimExpr], predicate: PrimExpr | None = None, span: Span | None = None)

Buffer 存储节点。

参数:

buffer (Buffer) – 缓冲区。
value (PrimExpr) – 要存储的值。
indices (List[PrimExpr]) – 要存储的索引位置。
predicate (Optional[PrimExpr]) – 一个布尔值向量掩码，指示要存储的向量的哪些通道。掩码的通道数必须等于值中的通道数。
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.BufferRealize(buffer: Buffer, bounds: List[Range], condition: PrimExpr, body: Stmt, span: Span | None = None)

Buffer 实现节点。

参数:

buffer (Buffer) – 缓冲区。
bounds (List[Range]) – 要存储的值。
condition (PrimExpr) – 实现条件。
body (Stmt) – 语句的主体。
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.ProducerStore(producer: DataProducer, value: PrimExpr, indices: List[PrimExpr], span: Span | None = None)

ProducerStore 节点。

参数:

producer (DataProducer) – 数据生产者。
value (PrimExpr) – 要存储的值。
indices (list of Expr) – 存储的索引参数。
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.Allocate(buffer_var: Var, dtype: str, extents: List[PrimExpr], condition: PrimExpr, body: Stmt, annotations: Mapping[str, Object] | None = None, span: Span | None = None)

Allocate 节点。

参数:

buffer_var (tir.Var) – 缓冲区变量。
dtype (str) – 缓冲区的数据类型。
extents (list of Expr) – allocate 的范围
condition (PrimExpr) – 条件。
body (Stmt) – 主体语句。
annotations (Optional[Mapping[str, Object]]) – 附加的注解提示
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.AllocateConst(buffer_var: Var, dtype: str, extents: List[PrimExpr], data_or_idx: NDArray | int, body: Stmt, annotations: Mapping[str, Object] | None = None, span: Span | None = None)

分配常量节点。

参数:

buffer_var (tir.Var) – 缓冲区变量。
dtype (str) – 缓冲区的数据类型。
extents (list of Expr) – allocate 的范围
data_or_idx (Union[NDArray, int]) – 如果是 NDArray，这是与常量关联的常量数据。如果是整数，这是指 IRModule 的 “constants” 属性的索引，该属性包含 AllocateConst。
body (Stmt) – 主体语句。
annotations (Optional[Mapping[str, Object]]) – 关于分配的附加注解。
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.AttrStmt(node: Object, attr_key: str, value: PrimExpr, body: Stmt, span: Span | None = None)

AttrStmt 节点。

参数:

node (Object) – 要注解属性的节点
attr_key (str) – 属性类型键。
value (PrimExpr) – 属性的值
body (Stmt) – 主体语句。
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.DeclBuffer(buffer: Buffer, body: Stmt, span: Span | None = None)

DeclBuffer 节点。

参数:

buffer (Buffer) – 被声明的缓冲区。
body (Stmt) – 要执行的主体语句。
span (Optional[Span]) – 此 DeclBuffer 在源代码中的位置。

class tvm.tir.ProducerRealize(producer: DataProducer, bounds: List[Range], condition: PrimExpr, body: Stmt, storage_scope: str = '', span: Span | None = None)

ProducerRealize 节点。

参数:

producer (DataProducer) – 数据生产者。
bounds (List[Range]) – realize 的边界
condition (PrimExpr) – 实现条件。
body (Stmt) – realize 主体
storage_scope (str) – 与此 realization 关联的存储作用域
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.SeqStmt(seq: List[Stmt], span: Span | None = None)

语句序列。

参数:

seq (List[Stmt]) – 语句列表
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.IfThenElse(condition: PrimExpr, then_case: Stmt, else_case: Stmt | None, span: Span | None = None)

IfThenElse 节点。

参数:

condition (PrimExpr) – 条件表达式
then_case (Stmt) – 如果条件为真，则执行的语句。
else_case (Optional[Stmt]) – 如果条件为假，则执行的语句。
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.Evaluate(value: PrimExpr, span: Span | None = None)

Evaluate 节点。

参数:

value (PrimExpr) – 要计算的表达式。
span (Optional[Span]) – 语句在源代码中的位置。

class tvm.tir.Prefetch(buffer: Buffer, bounds: List[Range], span: Span | None = None)

Prefetch 节点。

参数:

buffer (Buffer) – 要预取的缓冲区。
bounds (List[Range]) – 要预取的边界。
span (Optional[Span]) – 语句在源代码中的位置。

tvm.tir.stmt_seq(*args: PrimExpr | Stmt) → SeqStmt

创建语句序列

参数:: *args (Union[PrimExpr, Stmt]) – 要组合为序列的语句列表。
返回值:: stmt – 组合后的语句。
返回类型:: Stmt

tvm.tir.stmt_list(stmt: Stmt) → List[Stmt]

从块创建语句列表。

参数:: stmt (Stmt) – 输入语句。
返回值:: stmt_list – 解包后的语句列表
返回类型:: List[Stmt]

class tvm.tir.BufferRegion(buffer: Buffer, region: List[Range])

BufferRegion 节点。

参数:

buffer (Buffer) – 缓冲区区域的缓冲区
region (List[Range]) – 缓冲区区域的区域数组

class tvm.tir.MatchBufferRegion(buffer: Buffer, source: BufferRegion)

MatchBufferRegion 节点。

参数:

buffer (Buffer) – 目标缓冲区
source (BufferRegion) – 源缓冲区的区域

class tvm.tir.Block(iter_vars: List[IterVar], reads: List[BufferRegion], writes: List[BufferRegion], name_hint: str, body: Stmt, init: Stmt | None = None, alloc_buffers: List[Buffer] | None = None, match_buffers: List[MatchBufferRegion] | None = None, annotations: Mapping[str, Object] | None = None, span: Span | None = None)

Block 节点。

参数:

iter_vars (List[IterVar]) – 块变量。
reads (List[BufferRegion]) – 块的读取缓冲区区域。
writes (List[BufferRegion]) – 块的写入缓冲区区域。
name_hint (str) – 块的 name_hint。
body (Stmt) – 块的主体。
init (Optional[Stmt]) – reduction 块的 init 块
alloc_buffers (Optional[list[Buffer]]) – 缓冲区分配
match_buffers (Optional[List[MatchBufferRegion]]) – 子区域缓冲区匹配
annotations (Optional[Mapping[str, Object]]) – 附加的注解提示。
span (Optional[Span]) – 此块在源代码中的位置。

class tvm.tir.BlockRealize(iter_values: List[PrimExpr], predicate: PrimExpr | bool, block: Block, span: Span | None = None)

BlockRealize 节点。

参数:

iter_values (List[PrimExpr]) – 块变量的绑定值。
predicate (Union[PrimExpr, bool]) – 块的谓词。
block (Block) – 要 realize 的块
span (Optional[Span]) – 此 block_realize 在源代码中的位置。

class tvm.tir.PrimFunc(params, body, ret_type=None, buffer_map=None, attrs=None, span=None)

函数声明表达式。

参数:

params (List[Union[tvm.tir.Var, tvm.tir.Buffer]]) – 函数的输入参数列表。
body (tvm.tir.Stmt) – 函数的主体。
ret_type (tvm.ir.Type) – 函数的返回类型注解。
buffer_map (Map[tvm.tir.Var, tvm.tir.Buffer]) – 缓冲区绑定映射。
attrs (Optional[tvm.Attrs]) – 函数的属性，可以为 None
span (Optional[Span]) – 此 itervar 在源代码中的位置。

with_body(new_body, span=None)

创建一个具有相同签名但新主体的 PrimFunc。

参数:

new_body (Stmt) – 新的主体。
span (Optional[Span]) – 此 itervar 在源代码中的位置。

返回值:

new_func – 创建的新函数。

返回类型:

PrimFunc

specialize(param_map: Mapping[Var, PrimExpr | Buffer])

特化 PrimFunc 的参数

参数:: param_map (Mapping[tir.Var, Union[PrimExpr, Buffer]]) – 从函数参数到实例的映射

示例

我们可以使用符号形状定义 Meta TIR 函数

@T.prim_func
def mem_copy(a: T.handle, b: T.handle, m: T.int32, n: T.int32) -> None:
    A = T.match_buffer(a, (m, n), "float32")
    B = T.match_buffer(b, (m, n), "float32")

    for i, j in T.grid(m, n):
        with T.block():
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj]

然后我们可以使用给定的形状或缓冲区使其特化。

a, _, m, n = mem_copy.params
func = mem_copy.specialize({a: tir.decl_buffer((16, 16))})
# or
func = mem_copy.specialize({n: 16, m: 16})

特化后的函数

@T.prim_func
def mem_copy_16_16(a: T.handle, b: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    B = T.match_buffer(b, (16, 16), "float32")

    for i, j in T.grid(16, 16):
        with T.block():
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj]

返回值:: func – 参数特化后的新函数
返回类型:: PrimFunc

class tvm.tir.TensorIntrin(desc, impl)

张量内联函数。

参数:

desc (PrimFunc) – 描述计算的函数。
impl (PrimFunc) – 执行实现的函数。

static register(name: str, desc: PrimFunc, impl: PrimFunc, override: bool = False)

使用名称注册张量内在函数。

参数:

name (str) – 要注册的 TensorIntrin 的名称。
desc (PrimFunc) – 描述计算的函数。
impl (PrimFunc) – 执行实现的函数。
override (bool) – 是否覆盖现有的内在函数。

static get(name: str, allow_missing: bool = False) → TensorIntrin | None

通过名称查找张量内在函数。

参数:

name (str) – 要查找的 TensorIntrin 的名称。
allow_missing (bool) – 是否允许缺少张量内在函数。如果为 False，则在张量内在函数
exist. (不存在) 时引发错误。

返回值:

result – 具有指定名称的 TensorIntrin，如果未找到则为 None。

返回类型:

Optional[TensorIntrin]

class tvm.tir.IndexMap(initial_indices, final_indices, inverse_index_map)

从多维索引到另一组多维索引的映射

参数:

initial_indices (List[tir.Var]) – 表示重映射之前的索引的变量。
final_indices (List[PrimExpr]) – 定义重映射之后的索引的表达式。
inverse_index_map (Union[Callable, Optional[IndexMap]]) – 可选的预定义逆索引映射。当定义此项时，IndexMap::Inverse 将返回预定义的逆索引映射。否则，逆索引映射将即时计算。用户有责任确保预定义逆索引映射的正确性。

static from_func(mapping_function: Callable, ndim: int | None = None, inverse_index_map: Callable | IndexMap | None = None, *, index_dtype: str = 'int64')

从函数创建索引映射

参数:

mapping_function (Callable) – 从源索引映射到目标索引的函数。该函数应接受 tir.Var 参数，并返回 tir.PrimExpr 或 tir.PrimExpr 列表。返回 tir.PrimExpr 等同于返回包含该 tir.PrimExpr 的长度为 1 的列表。
ndim (Optional[int]) – 此转换应应用到的缓冲区的维数。如果 mapping_function 使用可变参数 *args，则必须指定 ndim。如果 mapping_function 不使用可变参数，则 ndim 是可选的。
inverse_index_map (Union[Callable, Optional[IndexMap]]) – 可选的预定义逆索引映射。当定义此项时，IndexMap::Inverse 将返回预定义的逆索引映射。否则，逆索引映射将即时计算。用户有责任确保预定义逆索引映射的正确性。

返回值:

index_map – 返回表示 mapping_function 的 IndexMap。

返回类型:

IndexMap

static from_func_with_separators(mapping_function: Callable, ndim: int | None = None, inverse_index_map: Callable | IndexMap | None = None, *, index_dtype: str = 'int64')

从函数创建索引映射

参数:

mapping_function (Callable) – 从源索引映射到目标索引的函数。该函数应接受 tir.Var 参数，并返回 tir.PrimExpr 或列表。返回列表的每个元素都应该是 tir.PrimExpr 或对象 IndexMap.AXIS_SEPARATOR。返回 tir.PrimExpr 等同于返回包含该 tir.PrimExpr 的长度为 1 的列表。
ndim (Optional[int]) – 此转换应应用到的缓冲区的维数。如果 mapping_function 使用可变参数 *args，则必须指定 ndim。如果 mapping_function 不使用可变参数，则 ndim 是可选的。
inverse_index_map (Union[Callable, Optional[IndexMap]]) – 可选的预定义逆索引映射。当定义此项时，IndexMap::Inverse 将返回预定义的逆索引映射。否则，逆索引映射将即时计算。用户有责任确保预定义逆索引映射的正确性。
index_dtype (str) – 用于映射函数中输入迭代器的默认索引数据类型。

返回值:

ret – 返回一个元组，其第一个元素是表示 mapping_function 的 IndexMap，第二个索引是 IndexMap.AXIS_SEPARATOR 出现的索引列表。

返回类型:

Tuple[IndexMap, List[int]]

is_equivalent_to(other_map: IndexMap) → bool

返回索引映射是否等效。

参数:: other_map (IndexMap) – 应与之进行比较的 IndexMap。
返回值:: is_equivalent – 如果两个映射表示相同的转换，则为 True，否则为 False
返回类型:: bool

map_indices(indices: List[PrimExpr]) → List[PrimExpr]

将索引映射应用于一组索引

参数:: indices (List[PrimExpr]) – 要映射的索引
返回值:: result – 映射后的索引
返回类型:: List[PrimExpr]

map_shape(shape: List[PrimExpr]) → List[PrimExpr]

将索引映射应用于缓冲区形状

参数:: shape (List[PrimExpr]) – 要映射的缓冲区形状
返回值:: result – 映射后的形状
返回类型:: List[PrimExpr]

map_ndarray(arr_src: NDArray) → NDArray

应用此索引映射来转换输入 NDArray 的布局

参数:: arr_src (runtime.NDArray) – 要转换的 NDArray
返回值:: arr_dst – 转换后的 NDArray
返回类型:: runtime.NDArray

inverse(shape: List[Range | PrimExpr]) → IndexMap

返回映射的逆

如果函数不是双射的，则抛出错误。

参数:: shape (List[Union[Range,PrimExpr]]) – 应在其中确定逆的区域。用于验证映射在此范围内是否是双射的。
返回值:: inverse – 逆
返回类型:: IndexMap

non_surjective_inverse(shape: List[Range | PrimExpr]) → Tuple[IndexMap, PrimExpr]

返回映射的逆

可以应用于引入填充的转换。

参数:: shape (List[Union[Range,PrimExpr]]) – 应在其中确定逆的区域。用于确定谓词。
返回值:: result – 逆，以及谓词，对于该谓词，逆映射到输入范围内的有效索引。
返回类型:: Tuple[IndexMap, PrimExpr]

示例

index_map = IndexMap.from_func(lambda i: [i//4, i%4])
inverse_map, predicate = index_map.non_surjective_inverse([14])
assert inverse_map.is_equivalent_to(IndexMap.from_func(lambda j,k: [4*j + k])
print(predicate) # Prints "(axis0==3) && (axis2 >= 2)"

tvm.tir.call_packed_lowered(*args, span=None)

Packed 调用的降级版本。packed 函数的参数可以是 Expr 或 Buffer。当呈现 Expr 时，参数是相应的 POD 类型。当参数是 Buffer 时，相应的 PackedFunc 将接收一个 TVMArrayHandle，其内容在回调期间有效。如果 PackedFunc 是 python 回调，则相应的参数是 NDArray。

参数:

args (list of Expr or Buffer.) – 位置参数。
span (Optional[Span]) – 此运算符在源代码中的位置。

返回值:

call – 调用表达式。

返回类型: