tvm.ir

所有 IR 变体的通用数据结构。

类

Attrs	属性节点，主要用于定义运算符的属性。
DictAttrs	字典属性。
EnvFunc	环境函数。
Node	所有 IR 节点的基类。
SourceName(name)	源位置的标识符。
Span(source_name, line, end_line, column, ...)	指定源程序中的位置。
SequentialSpan(spans)	源 span 的序列
Array	TVM 的 Array 容器。
Map	TVM 的 Map 容器。
BaseExpr	所有表达式的基类。
GlobalVar(name_hint[, type_annot])	IR 中的全局变量。
PrimExpr	所有原始表达式的基类。
Range(begin[, end, span])	表示 TVM 中的范围。
RelaxExpr	所有非原始表达式的基类。
BaseFunc	所有函数的基类。
CallingConv(value)	可能的调用约定类型。
GlobalInfo	可以出现在 IR 中的所有全局信息的基节点
DummyGlobalInfo()
VDevice([target, vdevice_id, memory_scope])
IRModule([functions, attrs, global_infos])	保存函数和类型定义的 IRModule。
Op()	IR 中的原始运算符。
FuncType(arg_types, ret_type)	函数类型。
PointerType(element_type[, storage_scope])	在底层 TIR 中使用的 PointerType。
PrimType(dtype)	底层 IR 中的原始数据类型
TupleType(fields)	元组值的类型。
Type	所有类型的基类。

函数

make_node(type_key, **kwargs)	通过类型键和字段创建一个新的 IR 节点
assert_structural_equal(lhs, rhs[, ...])	断言 lhs 和 rhs 在结构上彼此相等。
load_json(json_str)	从 json_str 加载 tvm 对象。
save_json(node)	将 tvm 对象保存为 json 字符串。
structural_equal(lhs, rhs[, map_free_vars])	检查 lhs 和 rhs 的结构相等性。
structural_hash(node[, map_free_vars])	计算节点的结构哈希值
register_intrin_lowering(op_name, target, *)	注册 Op lowering 函数
register_op_attr(op_name, attr_key[, value, ...])	按名称注册运算符的运算符属性。

class tvm.ir.Attrs

属性节点，主要用于定义运算符的属性。

由 Python 端注册的函数使用，例如 compute、schedule 和 alter_layout。Attrs 作为这些函数的第一个参数传递。

方法

list_field_info()	获取字段信息
keys()	获取属性中的名称列表。
get_int_tuple(key)	获取键的 Python int 元组
get_int(key)	获取键的 Python int 值
get_str(key)	获取键的 Python int 值

list_field_info()

获取字段信息

返回值:

infos – 字段信息列表

返回类型:

list of AttrFieldInfo

keys()

获取属性中的名称列表。

返回值:

keys – 键列表

返回类型:

list of str

get_int_tuple(key)

获取键的 Python int 元组

参数:

key (str)

返回值:

值

返回类型:

Tuple of int

get_int(key)

获取键的 Python int 值

参数:

key (str)

返回值:

值

返回类型:

int

get_str(key)

获取键的 Python int 值

参数:

key (str)

返回值:

值

返回类型:

int

class tvm.ir.DictAttrs

字典属性。

方法

keys()	获取属性中的名称列表。
get(key[, default])	获取具有默认值的元素。
items()	从 map 获取 items。

keys()

获取属性中的名称列表。

返回值:

keys – 键列表

返回类型:

list of str

get(key, default=None)

获取具有默认值的元素。

items()

从 map 获取 items。

tvm.ir.make_node(type_key, **kwargs)

通过类型键和字段创建一个新的 IR 节点

参数:

• type_key (str) – 节点的类型键。

• **kwargs (dict) – 节点的字段。

返回值:

node – 相应的 IR 节点

返回类型:

Node

注意

如果创建的节点是 AttrsNode 的实例，则创建器函数还将运行边界检查和默认值设置，如同 Attrs 所支持的那样。

示例

以下代码构造了一个 IntImm 对象

x = tvm.ir.make_node("IntImm", dtype="int32", value=10, span=None)
assert isinstance(x, tvm.tir.IntImm)
assert x.value == 10

class tvm.ir.EnvFunc

环境函数。

这是一个全局函数对象，可以通过其名称进行序列化。

方法

get(name)

获取静态 env 函数

static get(name)

获取静态 env 函数

参数:

name (str) – 函数的名称。

class tvm.ir.Node

所有 IR 节点的基类。

class tvm.ir.SourceName(name)

源位置的标识符。

参数:

name (str) – 源的名称。

class tvm.ir.Span(source_name, line, end_line, column, end_column)

指定源程序中的位置。

参数:

• source (SourceName) – 源名称。

• lineno (int) – 行号。

• col_offset (int) – 位置的列偏移。

class tvm.ir.SequentialSpan(spans)

源 span 的序列

此 span 专门用于表达式，该表达式来自 IR 转换后的多个表达式。

参数:

spans (Array) – span 的数组。

tvm.ir.assert_structural_equal(lhs, rhs, map_free_vars=False)

断言 lhs 和 rhs 在结构上彼此相等。

参数:

• lhs (Object) – 左操作数。

• rhs (Object) – 左操作数。

• map_free_vars (bool) – 是否应将未绑定到任何定义的自由变量映射为彼此相等。

:raises ValueError : 如果断言不成立。

另请参阅

structural_equal

tvm.ir.load_json(json_str) → Object

从 json_str 加载 tvm 对象。

参数:

json_str (str) – json 字符串

返回值:

node – 加载的 tvm 节点。

返回类型:

Object

tvm.ir.save_json(node) → str

将 tvm 对象保存为 json 字符串。

参数:

node (Object) – 要保存的 TVM 对象。

返回值:

json_str – 已保存的 json 字符串。

返回类型:

str

tvm.ir.structural_equal(lhs, rhs, map_free_vars=False)

检查 lhs 和 rhs 的结构相等性。

结构相等性在 IRNode 的 DAG 中递归定义。有两种类型的节点

• 图节点：lhs 中的图节点只能映射为 rhs 中的一个且仅一个图节点。

• 普通节点：相等性是递归定义的，没有图节点的限制。

Vars(tir::Var, relax::Var) 是图节点。

如果满足以下条件之一，则 var 类型节点（例如 tir::Var）可以映射为与另一个相同类型的 var 相等

• 它们出现在相同的定义点（例如，函数参数）。

• 它们通过 same_as 关系指向相同的 VarNode。

• 它们出现在相同的使用点，并且 map_free_vars 设置为 True。

var 的规则用于重新映射函数参数和 let 绑定中出现的变量。

参数:

• lhs (Object) – 左操作数。

• rhs (Object) – 左操作数。

• map_free_vars (bool) – 自由变量（即没有定义位置的变量）是否应映射为彼此相等。

返回值:

result – 比较结果。

返回类型:

bool

另请参阅

structural_hash, assert_strucural_equal

tvm.ir.structural_hash(node, map_free_vars=False)

计算节点的结构哈希值

结构哈希值在 IRNode 的 DAG 中递归定义。有两种类型的节点

• 普通节点：哈希值仅由其内容和类型定义。

• 图节点：每个图节点将被分配一个唯一的索引，该索引按访问期间首次出现的顺序排列。图节点的哈希值由其内容的哈希值和索引组合而成。

structural_hash 被设计为与 structural_equal 一致。如果两个节点在结构上彼此相等，那么它们的结构哈希值（使用相同的 map_free_vars 选项）也应该彼此相等。

如果两个节点的结构哈希值彼此相等，那么很可能（除了罕见的哈希值冲突情况）这两个节点在结构上彼此相等。

参数:

• node (Object) – 要哈希的输入。

• map_free_vars (bool) – 如果 map_free_vars 设置为 true，我们将按自由变量出现的顺序对其进行哈希。否则，我们将按其内存指针地址进行哈希。

返回值:

result – 哈希结果

返回类型:

int

另请参阅

structrual_equal

class tvm.ir.Array

TVM 的 Array 容器。

您不需要显式创建 Array。通常，python 列表和元组在 tvm 函数调用期间会自动转换为 Array。您可能会在 TVM 函数调用的返回值中获得 Array。

class tvm.ir.Map

TVM 的 Map 容器。

您不需要显式创建 Map。通常，python 字典在 tvm 函数调用期间会自动转换为 Map。您可以使用 convert 创建一个 dict[Object-> Object] 到 Map 中

方法

items()	从 map 获取 items
get(key[, default])	获取具有默认值的元素。

items()

从 map 获取 items

get(key, default=None)

获取具有默认值的元素。

参数:

• key (object) – 属性键。

• default (object) – 默认对象。

返回值:

value – 结果值。

返回类型:

object

class tvm.ir.BaseExpr

所有表达式的基类。

class tvm.ir.GlobalVar(name_hint: str, type_annot: Type | None = None)

IR 中的全局变量。

GlobalVar 用于引用存储在 IRModule 中的全局函数。

参数:

name_hint (str) – 变量的名称。

class tvm.ir.PrimExpr

所有原始表达式的基类。

PrimExpr 用于底层代码优化和整数分析。

class tvm.ir.Range(begin: PrimExpr, end: PrimExpr | None = None, span: Span | None = None)

表示 TVM 中的范围。

您不需要显式创建 Range。在 API 函数中，Python 列表和元组将自动转换为 Range。

参数:

• begin (PrimExpr) – 当 end 为 None 时，范围的起始值。否则，它是范围的长度。

• end (Optional[PrimExpr]) – 范围的结束值。

• span (Optional[Span]) – 此节点在源代码中的位置。

注意

如果 end 参数不是 None，则构造函数创建范围 [begin, end)。否则，它创建 [0, begin)。

方法

from_min_extent(min_value, extent[, span])

通过最小值和范围长度构造 Range。

static from_min_extent(min_value: PrimExpr, extent: PrimExpr, span: Span | None = None) → Range

通过最小值和范围长度构造 Range。

这将构造一个范围在 [min_value, min_value + extent) 内的 Range。

参数:

• min_value (PrimExpr) – 范围的最小值。

• extent (PrimExpr) – 范围的长度。

• span (Optional[Span]) – 此节点在源代码中的位置。

返回值:

rng – 构造的范围。

返回类型:

Range

class tvm.ir.RelaxExpr

所有非原始表达式的基类。

属性

checked_type	获取 tvm.relax.Expr 的检查类型。
struct_info	获取结构信息字段

property checked_type

获取 tvm.relax.Expr 的检查类型。

返回值:

checked_type – 检查类型。

返回类型:

tvm.ir.Type

property struct_info: StructInfo | None

获取结构信息字段

返回值:

struct_info – 可用的结构信息。

返回类型:

tvm.relax.StructInfo

class tvm.ir.BaseFunc

所有函数的基类。

属性

attrs

返回函数的 attrs 成员。

方法

with_attr(attr_key_or_dict[, attr_value])	创建函数的新副本并更新属性。
with_attrs(attr_map)	复制 IRModule 并将给定的属性映射添加到其中。
without_attr(attr_key)	创建一个不包含指定键的属性的函数新副本。

property attrs

返回函数的 attrs 成员。

with_attr(attr_key_or_dict, attr_value=None) → BaseFunc

创建函数的新副本并更新属性。

参数:

• attr_key_or_dict (Union[str, dict]) – 要使用的属性键或包含多个键值对的字典。

• attr_value (Object) – 新的属性值。

返回值:

func – 函数的新副本

返回类型:

BaseFunc

with_attrs(attr_map: DictAttrs | Dict[str, Object]) → BaseFunc

复制 IRModule 并将给定的属性映射添加到其中。 :param attr_map: 属性映射 :type attr_map: Union[DictAttrs, Dict[str, Object]]

返回值:

func – 函数的新副本

返回类型:

BaseFunc

without_attr(attr_key: str) → BaseFunc

创建一个不包含指定键的属性的函数新副本。

参数:

attr_key (str) – 要从属性对中删除的属性键。

返回值:

func – 函数的新副本

返回类型:

BaseFunc

class tvm.ir.CallingConv(value)

可能的调用约定类型。

class tvm.ir.GlobalInfo

可以出现在 IR 中的所有全局信息的基节点

方法

same_as(other)

与结构相等性重载。

same_as(other)

与结构相等性重载。

class tvm.ir.DummyGlobalInfo

class tvm.ir.VDevice(target=None, vdevice_id: int = 0, memory_scope: str = 'global')

class tvm.ir.IRModule(functions=None, attrs=None, global_infos=None)

保存函数和类型定义的 IRModule。

IRModule 是跨堆栈的所有 IR 转换的基本单元。

参数:

functions (Optional[dict].) – 从全局变量到 BaseFunc 的映射

方法

functions_items()	按字母顺序获取 self.functions.items() 中的项。
update(other)	将另一个模块中的函数插入到当前模块中。
update_func(var, func)	更新模块中与全局变量对应的函数。
update_global_info(name, global_info)	更新模块中的全局信息
get_global_var(name)	按名称获取函数中的全局变量。
get_global_vars()	收集在此模块中定义的所有全局变量。
replace_global_vars(replacements)	替换模块内的 GlobalVar 实例
from_expr(expr[, functions])	从独立的表达式构造模块。
get_attr(attr_key)	获取 IRModule 属性。
with_attr(attr_key, attr_value)	复制 IRModule 并向其添加属性。
without_attr(attr_key)	复制 IRModule 并删除属性键及其关联的值。
with_attrs(attr_map)	复制 IRModule 并将给定的属性映射添加到其中。

functions_items()

按字母顺序获取 self.functions.items() 中的项。

返回值:

items – 函数项。

返回类型:

List[Tuple[GlobalVar, Function]]

update(other)

将另一个模块中的函数插入到当前模块中。

参数:

other (IRModule) – 要合并到当前模块中的模块。

update_func(var, func)

更新模块中与全局变量对应的函数。

参数:

• var (GlobalVar) – 全局变量。

• func (tvm.ir.BaseFunc) – 要插入的函数。

update_global_info(name, global_info)

更新模块中的全局信息

参数:

• name (str) – 全局信息的名称。

• global_info (List[GlobalInfo]) – 要更新的全局信息。

get_global_var(name)

按名称获取函数中的全局变量。

参数:

name (str) – 全局变量的名称。

返回值:

global_var – 映射到 name 的全局变量。

返回类型:

GlobalVar

Raises:

tvm.error.TVMError if we cannot find corresponding global var. – 如果我们找不到对应的全局变量，则抛出 tvm.error.TVMError。

get_global_vars()

收集在此模块中定义的所有全局变量。

返回值:

global_vars – 全局变量数组。

返回类型:

Array[GlobalVar]

replace_global_vars(replacements: Dict[str | GlobalVar, str | GlobalVar]) → IRModule

替换模块内的 GlobalVar 实例

替换 IRModule 中的 GlobalVars。由于 IRModule 可能包含对 GlobalVar 的内部引用，无论是在 TIR 还是 Relax 中，因此每当替换或重命名 GlobalVar 时都应使用此方法。

参数:

replacements (Dict[Union[str, _expr.GlobalVar], Union[str, _expr.GlobalVar]]) – 一个字典，其中每个键是要替换的 GlobalVar，对应的值是要替换它的 GlobalVar。

返回值:

更新后的模块

返回类型:

IRModule

static from_expr(expr, functions=None)

从独立的表达式构造模块。

参数:

• expr (RelaxExpr) – 起始表达式

• global_funcs (Optional[dict]) – 从全局变量到函数定义的映射

返回值:

mod – 包含传递的定义的模块，其中 expr 被设置为入口点（如果需要，包装在函数中）

返回类型:

模块

get_attr(attr_key)

获取 IRModule 属性。

参数:

attr_key (str) – 属性键。

返回值:

attr_value – 属性值

返回类型:

Any

with_attr(attr_key, attr_value)

复制 IRModule 并向其添加属性。

参数:

• attr_key (str) – 属性键。

• attr_value (Object) – 新的属性值。

返回值:

mod – 具有该属性的 IRModule 的新副本

返回类型:

IRModule

without_attr(attr_key: str) → IRModule

复制 IRModule 并删除属性键及其关联的值。 :param attr_key: 属性键。 :type attr_key: str

返回值:

mod – 不包含该属性的 IRModule 的新副本

返回类型:

IRModule

with_attrs(attr_map: DictAttrs | Dict[str, Object]) → IRModule

复制 IRModule 并将给定的属性映射添加到其中。 :param attr_map: 属性映射 :type attr_map: Union[DictAttrs, Dict[str, Object]]

返回值:

mod – 具有该属性的 IRModule 的新副本

返回类型:

IRModule

class tvm.ir.Op

IR 中的原始运算符。

方法

get(op_name)	获取给定名称的 Op
get_attr(attr_name)	获取关于运算符的附加属性。
has_attr(attr_name)	检查运算符是否具有附加属性。
set_attr(attr_name, value[, plevel])	设置关于运算符的属性。
reset_attr(attr_name)	重置关于运算符的属性。
add_argument(name, type, description)	将参数信息添加到函数。
set_support_level(level)	设置 op 的支持级别。
set_num_inputs(n)	设置 op 的支持级别。
set_attrs_type_key(key)	设置 op 的属性类型键。
list_op_names()	列出 op 注册表中的所有 op 名称。

static get(op_name)

获取给定名称的 Op

参数:

op_name (str) – 运算符名称

返回值:

op – 对应名称的 op

返回类型:

Op

get_attr(attr_name)

获取关于运算符的附加属性。

参数:

attr_name (str) – 属性名称。

返回值:

value – 属性值

返回类型:

object

has_attr(attr_name)

检查运算符是否具有附加属性。

参数:

attr_name (str) – 属性名称。

返回值:

value – 运算符是否具有附加属性

返回类型:

bool

set_attr(attr_name, value, plevel=10)

设置关于运算符的属性。

参数:

• attr_name (str) – 属性名称

• value (object) – 属性值

• plevel (int) – 优先级

reset_attr(attr_name)

重置关于运算符的属性。

参数:

attr_name (str) – 属性名称

add_argument(name, type, description)

将参数信息添加到函数。

参数:

• name (str) – 参数名称。

• type (str) – 参数类型。

• description (str) – 参数描述。

set_support_level(level)

设置 op 的支持级别。

参数:

level (int) – 支持级别。

set_num_inputs(n)

设置 op 的支持级别。

参数:

n (int) – 输入数量。

set_attrs_type_key(key)

设置 op 的属性类型键。

参数:

key (str) – 类型键。

static list_op_names()

列出 op 注册表中的所有 op 名称。

返回值:

value – 注册的 op 名称

返回类型:

List[str]

tvm.ir.register_intrin_lowering(op_name, target, *, f=None, level=10)

注册 Op lowering 函数

参数:

• op_name (str) – op 名称

• target (str) – 给定 intrinsic lowering 函数的目标字符串

• f (function, optional) – 要注册的函数。

• level (int) – 优先级

返回值:

fregister – 如果未指定 f，则注册 op lowering 函数。

返回类型:

function

tvm.ir.register_op_attr(op_name, attr_key, value=None, level=10)

按名称注册运算符的运算符属性。

参数:

• op_name (str) – 运算符的名称

• attr_key (str) – 属性名称。

• value (object, optional) – 要设置的值

• level (int, optional) – 优先级

返回值:

fregister – 如果未指定值，则注册函数。

返回类型:

function

class tvm.ir.FuncType(arg_types, ret_type)

函数类型。

函数类型由类型参数列表（用于启用泛型函数的定义）、一组类型约束（我们暂时忽略）、一系列参数类型和一个返回类型组成。

参数:

• arg_types (List[tvm.ir.Type]) – 参数类型

• ret_type (tvm.ir.Type) – 返回类型。

class tvm.ir.PointerType(element_type, storage_scope='')

在底层 TIR 中使用的 PointerType。

参数:

• element_type (tvm.ir.Type) – 指针元素的类型。

• storage_scope (str) – 指针寻址的存储范围。

class tvm.ir.PrimType(dtype)

底层 IR 中的原始数据类型

参数:

dtype (str) – 与 primtype 相关的运行时数据类型。

class tvm.ir.TupleType(fields)

元组值的类型。

参数:

fields (List[Type]) – 元组中的字段

class tvm.ir.Type

所有类型的基类。

方法

same_as(other)

通过引用相等性比较两个 TVM 类型。

same_as(other)

通过引用相等性比较两个 TVM 类型。