tvm.relax

包含 IR、类型、运算符、构建器、vm 等的 Relax IR 命名空间。

class tvm.relax.VirtualMachine(rt_mod: Module | Executable, device: Device | List[Device], memory_cfg: str | Dict[Device, str] | None = None, profile: bool = False)

Relax VM 运行时。

invoke_closure(closure: Object, *args: Any) → Object

调用闭包。

参数:

closure (Object) – VMClosure 对象。
args (list[tvm.runtime.NDArray] 或 list[np.ndarray]) – 闭包的参数。

返回:

result – 输出。

返回类型:

Object

save_function(func_name: str, saved_name: str, *args: List[Any], include_return: bool = True, **kwargs: Dict[str, Any]) → None

便捷函数。从模块中获取一个函数并保存一个 PackedFunc，当调用它时，将使用给定的参数调用该函数。PackedFunc 可以使用 saved_name 从模块访问。包含此函数是为了方便计时试验：调用返回的 PackedFunc 将比通常通过 VM 运行具有更少的字典查找开销。

如果已占用保存的名称，则可以覆盖它，但它不能覆盖 Relax 源代码中定义的函数的名称。

这实际上是在创建一个闭包，但是该函数具有不同的名称，以避免与 invoke_closure 混淆（它们并非旨在一起使用）。

参数:

func_name (str) – 应该打包的函数。
saved_name (str) – 结果闭包应保存的名称。
include_return (bool) – 保存的 PackedFunc 是否应返回其输出。如果通过 RPC 进行计时，则可能不希望在机器之间发送输出。
args (List[Any]) – 要与函数一起打包的参数。
kwargs (Dict[str, Any]) – 要与函数一起打包的任何命名参数

set_input(func_name: str, *args: Any, **kwargs: Any) → None

设置函数的输入。当通过 RPC 使用 VM 时，此接口通过在内部将参数中的 NDArray 转换为 DLTensor 来工作，DLTensor 在 RPC 中受支持，其中远程可能只有最小的 C 运行时。

注意：如果使用 set_input，则该函数必须使用 invoke_stateful 调用，并且必须使用 get_outputs 获取结果。

参数:

func_name (str) – 函数的名称。
args (List[tvm.runtime.NDArray] 或 List[np.ndarray]) – 函数的参数。
kwargs (dict of str to tvm.runtime.NDArray 或 np.ndarray) – 函数的命名参数。

invoke_stateful(func_name: str) → None

relax.使用 set_input 设置的参数从 VM 模块调用命名函数。在没有先使用 set_input 的情况下调用 invoke_stateful 是错误的（即使它是为了设置 0 个输入）；相反，如果已调用 set_input，则在不使用 invoke_stateful 的情况下调用该函数是错误的。

可以通过调用 get_outputs 获取调用的结果。

参数:: func_name (str) – 要调用的函数的名称。

get_outputs(func_name: str) → Object | Tuple[Any]

在调用 invoke_stateful 之后，获取给定名称的函数输出的值。

在没有先调用 invoke_stateful 的情况下调用此函数是错误的。

参数:: func_name (str) – 应获取其输出的函数的名称。
返回:: ret – 先前通过 invoke_stateful 调用函数的结果。如果结果是元组，则返回字段列表。字段也可能是元组，因此这些字段可以是任意嵌套的。
返回类型:: Union[tvm.Object, Tuple[Any]]

set_instrument(instrument: PackedFunc) → None

设置一个 instrumentation 函数。

如果存在 instrumentation，则将在每个 relax.Call 指令之前/之后调用该函数。该函数具有以下签名

def instrument(
    func: Union[VMClosure, PackedFunc],
    func_symbol: str,
    before_run: bool,
    ret_value: any,
    *args) -> bool:
    pass

instrumentation 接受以下参数： - func：要调用的函数对象。 - func_symbol：函数的符号名称。 - before_run：是在调用之前还是之后。 - ret_value：调用的返回值，仅在运行后有效。 - args：传递给调用的参数。

instrumentation 函数可以选择一个整数，该整数对应于后续运行的操作方向。有关更多详细信息，请参见 VMInstrumentReturnKind。

参数:: instrument (tvm.runtime.PackedFunc) – 一个 instrumentation 函数，每次 VM 调用指令都会调用它。

另请参阅

VMInstrumentReturnKind: VM 中可能的返回值。

time_evaluator(func_name: str, dev: Device, number: int = 10, repeat: int = 1, min_repeat_ms: int = 0, cooldown_interval_ms: int = 0, repeats_to_cooldown: int = 1, f_preproc: str = '') → Callable[[...], BenchmarkResult]

返回一个评估器，该评估器对模块中的函数进行计时。这遵循与 tvm.runtime.module 中的 time_evaluator 相同的约定。这可以与 save_function() 结合使用，以便计时避免额外的字典查找。

参数:

func_name (str) – 模块中函数的名称。
dev (Device) – 我们应该在其上运行此函数的设备。
number (int) – 运行此函数以取平均值的次数。我们将这些运行称为一次测量 repeat。
repeat (int, optional) – 重复测量次数。总共，该函数将被调用 (1 + number x repeat) 次，其中第一次是预热并将被丢弃。返回的结果包含 repeat 成本，每个成本都是 number 成本的平均值。
min_repeat_ms (int, optional) – 一次 repeat 的最小持续时间，以毫秒为单位。默认情况下，一次 repeat 包含 number 次运行。如果设置此参数，则将动态调整参数 number 以满足一次 repeat 的最小持续时间要求。即，当一次 repeat 的运行时间低于此时间时，number 参数将自动增加。
cooldown_interval_ms (int, optional) – 由 repeats_to_cooldown 定义的重复次数之间的冷却间隔，以毫秒为单位。
repeats_to_cooldown (int, optional) – 冷却激活之前的重复次数。
f_preproc (str, optional) – 我们希望在执行时间评估器之前执行的预处理函数名称。

注意

该函数将被调用 (1 + number x repeat) 次，第一次调用将被丢弃，以防出现延迟初始化。

示例

VM 函数的正常用法（如果该函数返回元组，则可能无法通过 RPC 工作）

target = tvm.target.Target("llvm", host="llvm")
ex = tvm.compile(TestTimeEvaluator, target)
vm = relax.VirtualMachine(mod, tvm.cpu())
timing_res = vm.time_evaluator("func_name", tvm.cpu())(arg0, arg1, ..., argn)

与有状态 API 一起使用

target = tvm.target.Target("llvm", host="llvm")
ex = tvm.compile(TestTimeEvaluator, target)
vm = relax.VirtualMachine(mod, tvm.cpu())
vm.set_input("func_name", arg0, arg1, ..., argn)
timing_res = vm.time_evaluator("invoke_stateful", tvm.cpu())("func_name")

通过 save_function 保存的闭包（这会导致在定时部分中减少字典查找）

target = tvm.target.Target("llvm", host="llvm")
ex = tvm.compile(TestTimeEvaluator, target)
vm = relax.VirtualMachine(mod, tvm.cpu())
vm.save_function("func_name", "func_name_saved", arg0, arg1, ..., argn)
timing_res = vm.time_evaluator("func_name_saved", tvm.cpu())()

返回:: ftimer – 该函数接受与 func 相同的参数，并返回 BenchmarkResult。ProfileResult 报告 repeat 时间成本（以秒为单位）。
返回类型:: function

profile(func_name: str, *args)

分析函数调用。

参数:

func_name (str) – 函数的名称。
args (List of NDArray 或 PackedFunc 支持的其他对象。) – 函数的参数。

返回:

report – 格式化的分析结果，显示每个操作的计时测量。

返回类型:

tvm.runtime.profiling.Report

class tvm.relax.VMInstrumentReturnKind(value): 一个枚举。

tvm.relax.Expr: 别名 RelaxExpr

class tvm.relax.Id: Var 中使用的唯一标识符（名称）。保证在所有 pass 中保持稳定。

class tvm.relax.Var(name_hint: str | Id, struct_info: StructInfo | None = None, span: Span | None = None)

所有 Relax 绑定的变量类。

参数:

name_hint (Union[str, Id]) – 变量的名称提示。
struct_info (Optional[StructInfo]) – 变量的 struct info 注解。
span (Optional[Span]) – 指向原始源代码的 Span

property name_hint: str: 获取当前 var 的名称提示。

class tvm.relax.DataflowVar(name_hint: str | Id, struct_info: StructInfo | None = None, span: Span | None = None)

变量节点的子类型，用于标记来自普通可见“函数局部”绑定的数据流变量。

参数:

name_hint (Union[str, Id]) – 变量的名称提示。
struct_info (Optional[StructInfo]) – 变量的 struct info 注解。
span (Optional[Span]) – 指向原始源代码的 Span

class tvm.relax.Binding: Relax 中绑定的基类。

class tvm.relax.MatchCast(var: Var, value: RelaxExpr, struct_info: StructInfo, span: Span | None = None)

运行时将值匹配到结构信息。

此操作执行运行时检查，在首次出现时填充未定义的符号形状变量和 struct_info 中的变量，并在其他情况下插入相等性断言。

参数:

var (relax.Var) – 匹配转换绑定到的返回变量。
value (Expr) – 输入值表达式。
struct_info (tvm.relax.StructInfo) – 要匹配转换到的结构信息。

class tvm.relax.VarBinding(var: Var, value: RelaxExpr, span: Span | None = None)

变量绑定，将左侧的变量与右侧绑定。

参数:

var (relax.Var) – 匹配转换绑定到的返回变量。
value (Expr) – 输入值表达式。

class tvm.relax.BindingBlock(bindings: List[Binding], span: Span | None = None): 绑定块的基类，内部的绑定可能是不纯的（带有副作用或控制流）。

class tvm.relax.DataflowBlock(bindings: List[Binding], span: Span | None = None): 数据流块，内部的绑定是纯的（没有副作用和控制流）。

class tvm.relax.SeqExpr(blocks: List[BindingBlock], body: RelaxExpr, span: Span | None = None): 由一系列绑定块和随后的表达式组成的序列。

class tvm.relax.ShapeExpr(values: List[PrimExpr] | Tuple[PrimExpr, ...] | Array, span: Span | None = None)

形状表达式，允许用户构造包含 PrimExpr 的形状。

参数:

values (Union[List[PrimExpr], Tuple[PrimExpr, ...], tvm.ir.Array]) – 形状表达式的值。
span (Optional[Span]) – 指向原始源代码的 Span

class tvm.relax.Tuple(fields: List[RelaxExpr] | Tuple[RelaxExpr, ...], span: Span | None = None)

将多个字段组合在一起的元组表达式。

参数:

fields (Union[List[Expr], Tuple[Expr, ...]]) – 元组中的字段。
span (Optional[Span]) – 指向原始源代码的 Span

class tvm.relax.TupleGetItem(tuple_value: RelaxExpr, index: int, span: Span | None = None)

从元组中获取索引为 index 的项。

参数:

tuple_value (Expr) – 输入元组表达式。
index (int) – 索引。
span (Optional[Span]) – 指向原始源代码的 Span

class tvm.relax.Function(params: List[Var], body: RelaxExpr, ret_struct_info: StructInfo | None = None, is_pure: bool | None = True, attrs: DictAttrs | None = None, span: Span | None = None)

Relax 函数。

static create_empty(params: List[Var], ret_struct_info: StructInfo, is_pure: bool | None = True, attrs: DictAttrs | None = None, span: Span | None = None): 构造一个 relax.Function，但不包含 body。

bind_symbolic_vars(binding_map: Mapping[str | Var, PrimExpr]) → Function

返回一个更新了符号变量的新函数。

参数:: binding_map (Mapping[Union[str, tvm.tir.Var], PrimExpr]) – 要替换的值的映射。键可以是 tir.Var 或变量的字符串名称。如果变量通过名称引用，则名称必须唯一标识函数中的符号变量。
返回:: func – 更新后的函数
返回类型:: Function

返回一个更新了符号变量的新函数。

参数:

binding_map (Mapping[) – Union[str, relax.Var], Union[int, float, PrimExpr, tvm.runtime.NDArray, _np.ndarray, Expr],
] –
要替换的值的映射。

键可以是 relax.Var 或 Relax 变量的字符串名称。如果变量通过名称引用，则名称必须唯一标识函数中的参数。

值必须是 relax 表达式，或可转换为 relax 表达式的值。该值必须与要替换的变量兼容。

返回:

func – 更新后的函数

返回类型:

Function

class tvm.relax.ExternFunc(global_symbol: String, struct_info: StructInfo | None = None, span: Span | None = None): 外部函数，表示 PackedFunc。

Relax 中的函数调用节点。

relax.Call 节点对应于计算图术语中的运算符应用节点。

参数:

op (tvm.ir.Op 或任何具有函数类型的 tvm.relax.Expr。) – 要调用的操作。
args (Union[List[Expr], Tuple[Expr, ...]]) – 调用的参数。
attrs (Optional[tvm.ir.Attrs]) – 调用的属性，可以为 None。
sinfo_args (Optional[Union[List[StructInfo], Tuple[StructInfo, ...]]]) – CallNode 的结构信息参数。sinfo_args 旨在仅对内部操作（例如，call_tir、call_builtin_with_ctx 等）和对 ExternFuncs 的调用为非空，主要用于结构信息推断。
span (Optional[Span]) – 指向原始源代码的 Span

class tvm.relax.If(cond: RelaxExpr, true_branch: RelaxExpr, false_branch: RelaxExpr, span: Span | None = None)

Relax 中的条件表达式。

参数:

cond (Expr) – 条件。
true_branch (Expr) – 当条件为真时评估的表达式。
false_branch (Expr) – 当条件为假时评估的表达式。
span (Optional[Span]) – 指向原始源代码的 Span

class tvm.relax.Constant(data: NDArray, struct_info: StructInfo | None = None, span: Span | None = None)

恒定张量

参数:

data (tvm.nd.NDArray) – 恒定张量的数据。
struct_info (Optional[StructInfo]) – 恒定张量的结构信息。如果未指定，则从数据推断。
span (Optional[Span]) – 指向原始源代码的 Span

注意

标量常量由 0 维常量张量表示。

class tvm.relax.PrimValue(value: PrimExpr | int, span: Span | None = None): 表示值的 prim expr。

class tvm.relax.DataTypeImm(value: DataType | str, span: Span | None = None): 表示数据类型常量。

class tvm.relax.StringImm(value: str, span: Span | None = None): 表示字符串字面量常量。

创建一个常量值。

参数:

value (Union[bool, int, float, numpy.ndarray, tvm.nd.NDArray]) – 常量值。
dtype (Optional[str]) – 结果常量的数据类型。

注意

当 dtype 为 None 时，我们使用以下规则

int 映射到 “int32”
float 映射到 “float32”
bool 映射到 “bool”
其他类型使用与 numpy 相同的默认规则。

tvm.relax.extern(name: str, struct_info: StructInfo | None = None, span: Span | None = None): 创建外部函数。

tvm.relax.get_shape_of(expr: RelaxExpr) → RelaxExpr

获取 expr 的形状。

参数:: expr (Expr) – 输入 expr。
返回:: shape – 形状表达式
返回类型:: Expr

注意

此函数要求 expr 是归一化的。如果 expr 的 StructInfo 不是 TensorStructInfo，该函数将报告错误。它会尝试在可能的情况下返回符号函数。如果张量没有编译时符号形状，则该函数将选择返回 relax.Call(relax.op.shape_of, [expr])。

class tvm.relax.ObjectType(span: Span | None = None): 一种对应于 tvm::runtime::Object 的类型，是 TVM 中所有可能的对象值的基类。

class tvm.relax.ShapeType(ndim: int = -1, span: Span | None = None)

Relax 中形状的类型。

参数:: ndim (Optional[int]) – 形状的大小。

class tvm.relax.TensorType(ndim=-1, dtype='float32', span: Span | None = None)

Relax 中的动态张量类型。

这是分配给具有已知 dtype 和未知形状的张量的类型。

参数:

ndim (Optional[int]) – 张量的维度。
dtype (Optional[str]) – 内容数据类型。

class tvm.relax.PackedFuncType(span: Span | None = None): Relax 中 ExternFunc 的类型。

class tvm.relax.ExecBuilder

一个用于发出指令并为虚拟机构建可执行文件的构建器。

r(idx: int) → int: 将指令的参数设置为寄存器。

imm(value: int) → int: 将指令的参数设置为立即数。

c(idx: int) → int: 将指令的参数设置为常量。

f(name: str) → int: 将指令的参数设置为函数。

declare_function(func_name: str, kind: VMFuncKind = VMFuncKind.PACKED_FUNC) → None: 声明一个函数

function(func_name: str, num_inputs: int | None = 0, param_names: List[str] | None = None) → VMFuncScope: 注释一个 VM 函数。

emit_call(name: str, args: List[NDArray | DataType] | None = None, dst: int | None = None) → None: 发出调用 packed function 的 call 指令。

emit_ret(result: int) → None: 发出 return 指令

emit_goto(pc_offset): 发出 goto 指令

emit_if(cond, false_offset): 发出 if 指令

get() → VMExecutable: 返回可执行文件

tvm.relax.call_tir(gvar: GlobalVar, args: RelaxExpr, out_sinfo: TensorStructInfo | List[TensorStructInfo], tir_vars: ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None = None) → Call

relax.Call 调用 tir.prim_func 并返回输出。

参数:

gvar (GlobalVar) – 指向 tir PrimFunc 的 GlobalVar。
args (Expr) – 输入参数。
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – call_tir 输出的结构信息。它应该是单个或 TensorStructInfo 列表。每个都表示返回张量的结构信息。
tir_vars (Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]) – ShapeExpr，表示调用 func 时要解包的整数元组。如果未使用则为空

返回:

ret – call_tir 运算符的调用节点。

返回类型:

relax.Call

relax.调用 TIR PrimFunc 并返回结果，执行指定的原地计算（基于 inplace_indices 参数；输出将别名化为原地索引选择的输入）。

警告：此运算符被类型系统认为是纯的，但实际上会改变 inplace_indices 指定的参数。此运算符不应直接使用，而应由已检查过执行原地操作是否安全的 pass 插入（即，指定为输出的参数均未被别名化，或在调用 call_tir_inplace 后仍然存活）。

直接调用此运算符仅应用于测试目的。

参数:

gvar (GlobalVar) – 指向 TIR PrimFunc 的 GlobalVar。
args (Expr) – 输入参数。
inplace_indices (Union[int, List[int]]) – 指定哪些参数应用于原地计算。如果 inplace_indices 是单个整数，它将被转换为单例列表。假设 inplace_indices[i] = j，其中 j >= 0。那么第 i 个输出将是 args[j] 的别名。如果 inplace_indices[i] = -1，那么第 i 个输出将是新分配的张量。 inplace_indices 中至少有一个成员不能为 -1。
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – call_tir_inplace 输出的结构信息。它应该是单个 TensorStructInfo 或 TensorStructInfo 列表。每个都表示返回张量的结构信息。如果给定 TensorStructInfo 列表，则结果将是 TensorStructInfo 元组。
tir_vars (Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]) – ShapeExpr，表示调用 func 时要解包的整数元组。如果未使用则为空

返回:

ret – call_tir 运算符的调用节点。

返回类型:

relax.Call

tvm.relax.call_pure_packed(func: str | ExternFunc | GlobalVar, *args: RelaxExpr, sinfo_args: StructInfo | List[StructInfo]) → RelaxExpr

构造对 packed 函数的调用，该函数应被视为纯函数，即使 packed 调用通常不被视为纯函数。

结果调用将具有与直接调用 packed 函数相同的语义。

注意：这应该用于用户知道用这些参数调用 packed 函数实际上不会导致任何副作用的情况。如果它用于确实导致副作用的调用，那么编译器最终可能会删除、重新排序或重复该调用，并且不对被调用者的任何副作用做出任何保证。

参数:

func (Union[str, ExternFunc]) – PackedFunc 或 ExternFunc 节点的名称（全局符号）。
args (Expr) – PackedFunc 的参数。
sinfo_args (Union[StructInfo, List[StructInfo]]) – 结构信息参数列表（给出返回值的结构信息）。

返回:

result – Relax 调用，对应于 call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), sinfo_args)

返回类型:

Expr

tvm.relax.call_dps_packed(func: str | RelaxExpr, args: RelaxExpr, out_sinfo: TensorStructInfo | List[TensorStructInfo]) → Call

relax.调用目标传递风格的 packed 函数并返回输出。

注意：假定调用的函数是 _纯函数_（除了修改指定输出参数之外）。如果函数 _确实_ 导致其他副作用，那么编译器最终可能会删除、重新排序或重复这些效果——无法保证。

参数:

func (Union[str, Expr]) – 目标传递风格函数，可以是 ExternFunc。
args (Expr) – 输入参数。
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – call_dps_packed 输出的结构信息。它应该是单个或 TensorStructInfo 列表。每个都表示返回张量的结构信息。

返回:

ret – call_dps_packed 运算符的调用节点。

返回类型:

relax.Call

tvm.relax.call_tir_with_grad(gvar: GlobalVar, args: RelaxExpr, out_sinfo: TensorStructInfo | List[TensorStructInfo], te_grad_name: str, te_grad_kwargs: Dict[str, Object] = None, tir_vars: ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None = None) → Call

relax.调用 tir.prim_func 并返回输出。此内联函数会将一个 te 梯度函数（由 te_grad_name 引用）绑定到 call_tir_with_grad 节点。te 梯度函数将由 Gradient pass 调用。

参数:

gvar (GlobalVar) – 指向 tir PrimFunc 的 GlobalVar。
args (Expr) – 输入参数。
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – call_tir_with_grad 输出的结构信息。它应该是单个或 TensorStructInfo 列表。每个都表示返回张量的结构信息。
te_grad_name (str) – 与 call_tir_with_grad 节点关联的 te 梯度函数的注册名称。必须作为关键字参数提供。
te_grad_kwargs (Dict[str, Object], optional) – 传递给 te 梯度函数的关键字参数。可选择作为关键字参数提供。默认值：{}。
tir_vars (Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]) – ShapeExpr，表示调用 func 时要解包的整数元组。如果未使用则为空

返回:

ret – call_tir_with_grad 运算符的调用节点。

返回类型:

relax.Call

class tvm.relax.ExprFunctor

在 Expr 上定义的抽象访问器。定义了表达式的默认调度，并实现了 memoization。

visit_expr(expr: RelaxExpr) → RelaxExpr: 将访问器应用于表达式。

class tvm.relax.PyExprVisitor

一个抽象的 ExprVisitor，在 python 端具有自定义方法。这是用户面对的类，用于方法覆盖继承。 _tvm_metadata 描述了要继承的类（“cls”），用户可以覆盖的方法（“methods”）。

注意：@relax.expr_functor.visitor 是任何继承类的正确使用所必需的。

另请参阅：visitor, _PyExprVisitor

示例

@relax.expr_functor.visitor
def MyExprVisitor(PyExprVisitor):
    ...

visit_expr(expr: RelaxExpr) → None

Expr 的通用调度器。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr(const Expr& expr)。

参数:: expr (Expr) – 要访问的 expr。

visit_binding(binding: Binding) → None

Binding 的通用调度器。用户可以自定义此函数以覆盖 C++ 端的 VisitBinding(const Binding& binding)。

参数:: binding (Binding) – 要访问的 binding。

visit_binding_block(block: BindingBlock) → None

BindingBlock 的通用调度器。用户可以自定义此函数以覆盖 C++ 端的 VisitBindingBlock(const BindingBlock& block)。

参数:: block (BindingBlock) – 要访问的 block。

visit_var_def(var: Var) → None

用于访问 var 定义站点的通用调度器。用户可以自定义此函数以覆盖 C++ 端的 VisitVarDef(const relax.Var& var)。请注意，visit_var_() 仅访问 Var 的使用站点。

参数:: var (relax.Var) – 要访问的 var。

visit_constant_(op: Constant) → None

访问 Constant。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const ConstantNode* op)。

参数:: op (Constant) – 要访问的 Constant。

visit_tuple_(op: Tuple) → None

访问 Tuple。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const TupleNode* op)。

参数:: op (Tuple) – 要访问的 Tuple。

visit_var_(op: Var) → None

访问 Var。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const VarNode* op)。

参数:: op (relax.Var) – 要访问的 relax.Var。

visit_dataflow_var_(op: DataflowVar) → None

访问 DataflowVar。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const DataflowVarNode* op)。

参数:: op (DataflowVar) – 要访问的 DataflowVar。

visit_shape_expr_(op: ShapeExpr) → None

访问 ShapeExpr。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const ShapeExprNode* op)。

参数:: op (ShapeExpr) – 要访问的 ShapeExpr。

visit_extern_func_(op: ExternFunc) → None

访问 ExternFunc。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const ExternFuncNode* op)。

参数:: op (ExternFunc) – 要访问的 ExternFunc。

visit_global_var_(op: GlobalVar) → None

访问 GlobalVar。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const GlobalVarNode* op)。

参数:: op (GlobalVar) – 要访问的 GlobalVar。

visit_function_(op: Function) → None

访问 Function。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const FunctionNode* op)。

参数:: op (Function) – 要访问的 Function。

visit_call_(op: Call) → None

访问 Call。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const CallNode* op)。

参数:: op (relax.Call) – 要访问的 relax.Call。

visit_seq_expr_(op: SeqExpr) → None

访问 SeqExpr。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const SeqExprNode* op)。

参数:: op (SeqExpr) – 要访问的 SeqExpr。

visit_if_(op: If) → None

访问 If。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const IfNode* op)。

参数:: op (If) – 要访问的 If。

visit_op_(op: Op) → None

访问 Op。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const OpNode* op)。

参数:: op (Op) – 要访问的 Op。

visit_tuple_getitem_(op: TupleGetItem) → None

访问 TupleGetItem。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const TupleGetItemNode* op)。

参数:: op (TupleGetItem) – 要访问的 TupleGetItem。

visit_prim_value_(op: PrimValue) → None

访问 PrimValue。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const PrimValueNode* op)。

参数:: op (PrimValue) – 要访问的 PrimValue。

visit_string_imm_(op: StringImm) → None

访问 StringImm。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const StringImmNode* op)。

参数:: op (StringImm) – 要访问的 StringImm。

visit_data_type_imm_(op: DataTypeImm) → None

访问 DataTypeImm。用户可以自定义此函数，以在 C++ 端重写 VisitExpr_(const DataTypeImmNode* op)。

参数:: op (DataTypeImm) – 要访问的 DataTypeImm。

visit_var_binding_(binding: VarBinding) → None

访问 VarBinding。用户可以自定义此函数，以在 C++ 端重写 VisitBinding_(const VarBindingNode* binding)。

参数:: binding (VarBinding) – 要访问的 VarBinding。

visit_match_cast_(binding: MatchCast) → None

访问 MatchCast。用户可以自定义此函数，以在 C++ 端重写 VisitBinding_(const MatchCastNode* binding)。

参数:: binding (MatchCast) – 要访问的 MatchCast。

visit_binding_block_(block: BindingBlock) → None

访问 BindingBlock。用户可以自定义此函数，以在 C++ 端重写 VisitBindingBlock_(const BindingBlockNode* block)。

参数:: block (BindingBlock) – 要访问的 BindingBlock。

visit_dataflow_block_(block: DataflowBlock) → None

访问 DataflowBlock。用户可以自定义此函数，以在 C++ 端重写 VisitDataflowBlock_(const DataflowBlockNode* block)。

参数:: block (DataflowBlock) – 要访问的 DataflowBlock。

visit_var_def_(var: Var) → None

访问 relax.Var 定义站点。用户可以自定义此函数，以在 C++ 端重写 VisitVarDef_(const VarNode* var)。

参数:: var (relax.Var) – 要访问的 relax.Var。

visit_dataflow_var_def_(var: DataflowVar) → None

访问 DataflowVar 定义站点。用户可以自定义此函数，以在 C++ 端重写 VisitDataflowVarDef_(const DataflowVarNode* var)。

参数:: var (DataflowVar) – 要访问的 DataflowVar。

visit_span(span: Span) → None

访问 Span。用户可以自定义此函数，以在 C++ 端重写 VisitSpan(const Span& span)。

参数:: span (Span) – 要访问的 Span。

class tvm.relax.PyExprMutator(mod: IRModule | None = None)

一个抽象的 ExprMutator，在 python 端具有自定义方法。这是用户面对的用于方法重写继承的类。_tvm_metadata 描述了要继承的类（“cls”），用户可以重写的方法（“methods”），以及构造函数的参数（“fields”）。

注意：任何继承类要正确使用，都需要 @relax.expr_functor.mutator。

另请参阅：visitor, _PyExprVisitor

示例

@relax.expr_functor.mutator
def MyExprMutator(PyExprMutator):
    ...

visit_expr(expr: RelaxExpr) → RelaxExpr

Expr 的通用调度器。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr(const Expr& expr)。

参数:: expr (Expr) – 要访问的 expr。
返回:: result – 转换后的 Expr。
返回类型:: Expr

visit_binding(binding: Binding) → None

Binding 的通用调度器。用户可以自定义此函数以覆盖 C++ 端的 VisitBinding(const Binding& binding)。

参数:: binding (Binding) – 要访问的 binding。

visit_binding_block(block: BindingBlock) → BindingBlock

BindingBlock 的通用调度器。用户可以自定义此函数以覆盖 C++ 端的 VisitBindingBlock(const BindingBlock& block)。

参数:: block (BindingBlock) – 要访问的 block。
返回:: result – 转换后的 binding block。
返回类型:: BindingBlock

visit_var_def(var: Var) → Var

用于访问 var 定义站点的通用调度器。用户可以自定义此函数以覆盖 C++ 端的 VisitVarDef(const relax.Var& var)。请注意，visit_var_() 仅访问 Var 的使用站点。

参数:: var (relax.Var) – 要访问的 var。
返回:: result – 后序重写后的 var。
返回类型:: relax.Var

visit_constant_(op: Constant) → RelaxExpr

访问 Constant。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const ConstantNode* op)。

参数:: op (Constant) – 要访问的 Constant。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_tuple_(op: Tuple) → RelaxExpr

访问 Tuple。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const TupleNode* op)。

参数:: op (Tuple) – 要访问的 Tuple。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_var_(op: Var) → RelaxExpr

访问 Var。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const VarNode* op)。

参数:: op (relax.Var) – 要访问的 relax.Var。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_dataflow_var_(op: DataflowVar) → RelaxExpr

访问 DataflowVar。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const DataflowVarNode* op)。

参数:: op (DataflowVar) – 要访问的 DataflowVar。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_shape_expr_(op: ShapeExpr) → RelaxExpr

访问 ShapeExpr。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const ShapeExprNode* op)。

参数:: op (ShapeExpr) – 要访问的 ShapeExpr。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_extern_func_(op: ExternFunc) → RelaxExpr

访问 ExternFunc。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const ExternFuncNode* op)。

参数:: op (ExternFunc) – 要访问的 ExternFunc。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_global_var_(op: GlobalVar) → RelaxExpr

访问 GlobalVar。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const GlobalVarNode* op)。

参数:: op (GlobalVar) – 要访问的 GlobalVar。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_function_(op: Function) → RelaxExpr

访问 Function。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const FunctionNode* op)。

参数:: op (Function) – 要访问的 Function。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_call_(op: Call) → RelaxExpr

访问 Call。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const CallNode* op)。

参数:: op (relax.Call) – 要访问的 relax.Call。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_seq_expr_(op: SeqExpr) → RelaxExpr

访问 SeqExpr。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const SeqExprNode* op)。

参数:: op (SeqExpr) – 要访问的 SeqExpr。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_if_(op: If) → RelaxExpr

访问 If。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const IfNode* op)。

参数:: op (If) – 要访问的 If。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_op_(op: Op) → RelaxExpr

访问 Op。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const OpNode* op)。

参数:: op (Op) – 要访问的 Op。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_tuple_getitem_(op: TupleGetItem) → RelaxExpr

访问 TupleGetItem。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const TupleGetItemNode* op)。

参数:: op (TupleGetItem) – 要访问的 TupleGetItem。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_prim_value_(op: PrimValue) → RelaxExpr

访问 PrimValue。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const PrimValueNode* op)。

参数:: op (PrimValue) – 要访问的 PrimValue。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_string_imm_(op: StringImm) → RelaxExpr

访问 StringImm。用户可以自定义此函数以覆盖 C++ 端的 VisitExpr_(const StringImmNode* op)。

参数:: op (StringImm) – 要访问的 StringImm。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_data_type_imm_(op: DataTypeImm) → RelaxExpr

访问 DataTypeImm。用户可以自定义此函数，以在 C++ 端重写 VisitExpr_(const DataTypeImmNode* op)。

参数:: op (DataTypeImm) – 要访问的 DataTypeImm。
返回:: result – 转换后的 Expr
返回类型:: Expr

visit_var_binding_(binding: VarBinding) → None

访问 VarBinding。用户可以自定义此函数，以在 C++ 端重写 VisitBinding_(const VarBindingNode* binding)。

参数:: binding (VarBinding) – 要访问的 VarBinding。

visit_match_cast_(binding: MatchCast) → None

访问 MatchCast。用户可以自定义此函数，以在 C++ 端重写 VisitBinding_(const MatchCastNode* binding)。

参数:: binding (MatchCast) – 要访问的 MatchCast。

visit_binding_block_(block: BindingBlock) → BindingBlock

访问 BindingBlock。用户可以自定义此函数，以在 C++ 端重写 VisitBindingBlock_(const BindingBlockNode* block)。

参数:: block (BindingBlock) – 要访问的 BindingBlock。
返回:: result – 转换后的 binding block
返回类型:: BindingBlock

visit_dataflow_block_(block: DataflowBlock) → BindingBlock

访问 DataflowBlock。用户可以自定义此函数，以在 C++ 端重写 VisitDataflowBlock_(const DataflowBlockNode* block)。

参数:: block (DataflowBlock) – 要访问的 DataflowBlock。
返回:: result – 转换后的 binding block
返回类型:: BindingBlock

visit_var_def_(var: Var) → Var

访问 relax.Var 定义站点。用户可以自定义此函数，以在 C++ 端重写 VisitVarDef_(const VarNode* var)。

参数:: var (relax.Var) – 要访问的 relax.Var。
返回:: result – 后序重写后的 var。
返回类型:: relax.Var

visit_dataflow_var_def_(var: DataflowVar) → Var

访问 DataflowVar 定义站点。用户可以自定义此函数，以在 C++ 端重写 VisitDataflowVarDef_(const DataflowVarNode* var)。

参数:: var (DataflowVar) – 要访问的 DataflowVar。
返回:: result – 后序重写后的 var。
返回类型:: relax.Var

visit_span(span: Span) → Span

访问 Span。用户可以自定义此函数，以在 C++ 端重写 VisitSpan(const Span& span)。

参数:: span (Span) – 要访问的 Span。
返回:: result – 转换后的 span。
返回类型:: Span

visit_expr_post_order(expr: RelaxExpr) → RelaxExpr

后序重写一个 Expr 并进行规范化。

参数:: expr (Expr) – 要重写的 Expr。
返回:: result – 后序重写后的 Expr。
返回类型:: Expr

set_var_remap(vid: Id, var: Var) → None

在使用点将一个 var 重新映射到一个新的 var。

参数:

vid (Id) – 旧 var 的 vid。
var (relax.Var) – 新 var。

get_var_remap(vid: Id) → Var

在使用点将一个 var 重新映射到一个新的 var。

参数:: vid (Id) – 旧 var 的 vid
返回:: var – 重新映射的 var。
返回类型:: relax.Var

visit_with_new_scope(expr: RelaxExpr) → RelaxExpr

使用新的作用域重写 expr，用于 Function 的主体和 If 的分支中。

参数:: expr (Expr) – 要访问的 expr。
返回:: var – 访问后的 expr。
返回类型:: relax.Var

lookup_binding(var: Var) → RelaxExpr | None

查找绑定到变量的值。注意：对于函数参数，此函数返回 NullOpt。

参数:: var (relax.Var) – 要查找的 var。
返回:: var – 绑定到输入 var 的值。
返回类型:: relax.Var

with_struct_info(var: Var, struct_info: StructInfo) → Var

如果原始 var 的形状或类型与指定的形状或类型不匹配，则创建一个具有指定形状和类型的新 var。

参数:

var (relax.Var) – 要更新的 var。
struct_info (StructInfo) – struct info。

返回:

var – 填充了形状和类型的 var。

返回类型:

relax.Var

class tvm.relax.StructInfo

所有 StructInfo 的基类。

StructInfo 包含静态类型和运行时结构信息。

same_as(other): 重载结构相等性。

is_base_of(derived: StructInfo) → bool

检查 self 是否是另一个派生结构信息的基础。

参数:: derived (StructInfo) – 要检查的派生结构信息。
返回:: result – 检查结果。
返回类型:: bool

class tvm.relax.ObjectStructInfo(span: Span | None = None): Object 的结构信息。

原始 POD 值的结构信息。

参数:: dtype_or_expr (Union[str, DataType, PrimExpr]) – 原始值的数据类型，或原始值的已知表达式。

class tvm.relax.ShapeStructInfo(values: List[PrimExpr] | None = None, ndim: int = -1, span: Span | None = None)

形状值的结构信息。

参数:

values (Optional[List[PrimExpr]]) – 如果已知，则为符号形状值。
ndim (Optional[int]) – 形状的大小。

注意

请勿同时指定 values 和 ndim。

Tensor 值的结构信息。

参数:

shape (Optional[Expr]) – 形状表达式。
dtype (Optional[str]) – 内容数据类型。
vdevice (Optional[Vdevice]) – 虚拟设备。
ndim (Optional[int]) – tensor 的维度数量。

注意

请勿同时指定 shape 和 ndim。

class tvm.relax.TupleStructInfo(fields: List[StructInfo], span: Span | None = None)

Tuple 值的结构信息。

参数:: fields (List[StructInfo]) – 字段的结构信息。

class tvm.relax.FuncStructInfo(params: List[StructInfo], ret: StructInfo, purity: bool = True, span: Span | None = None)

函数值的结构信息。

参数:

params (List[StructInfo]) – 字段的结构信息。
ret (StructInfo) – 返回值的结构信息
purity (bool) – 函数是否是纯函数（没有可见的副作用）。注意：我们认为一个函数只有在所有输入上都是纯函数时才是纯函数。如果一个函数仅在某些情况下可能具有可见的副作用，我们仍然认为它是不纯的。

static opaque_func(*, ret: StructInfo | None = None, derive_func: str | EnvFunc | None = None, purity: bool = False, span: Span | None = None) → FuncStructInfo

创建一个不透明的 FuncStructInfo。

不透明函数接受一个指定返回值的结构信息的 ret，或者一个提供自定义推导规则的 derive_func。

参数:

ret (Optional[StructInfo]) – 函数返回值的结构信息。
derive_func (Optional[Union[str,EnvFunc]]) – 用于推导的环境函数
purity (bool) – 函数是否是纯函数（默认为 false，因为大多数不透明函数都不是纯函数）
span (Optional[Span]) – ast 的可选 span 信息。

返回:

info

返回类型:

FuncStructInfo

注意

我们不能同时指定 ret 和 derive_func。

tvm.relax.get_default_pipeline(target: Target): 获取给定目标的默认 Relax 编译 pipeline。

tvm.relax.get_pipeline(name: str = 'zero', **kwargs) → Pass

通过名称获取预构建的 pipeline

参数:

name (Optional[str]) – pipeline 的名称
kwargs (Dict[str, object]) – 用于配置 pipeline 的关键字参数。

返回:

pipeline – 转换 pipeline。

返回类型:

tvm.transform.Pass

tvm.relax.register_pipeline(name: str)

注册一个新的 pipeline

参数:: name (str) – pipeline 的名称

tvm.relax.convert_to_expr(value: Any) → RelaxExpr

辅助函数，用于将输入转换为 Expr，它遵循以下规则：1. 如果输入已经是 relax.Expr，则返回输入本身；2. 如果输入是 PrimExpr，则返回 relax.PrimValue；3. 如果输入是 tvm.String 或 str，则返回 relax.StringImm；4. 如果输入是 Expr 的 tuple/list，则返回 relax.Tuple。

注释

不允许使用 tvm.tir.StringImm，因为它具有歧义性，可能是 relax.StringImm 或 relax.PrimValue。

构建一个 IRModule 到 VM 可执行文件。

参数:

mod (IRModule) – 要构建的输入 IRModule。
target (Optional[Union[str, tvm.target.Target]]) –
一个构建目标，可以具有可选的主机端编译目标。

当 TVM 编译特定于设备的程序（如 CUDA）时，我们还需要主机（CPU）端代码与驱动程序交互，以正确设置维度和参数。host 用于指定主机端代码生成目标。默认情况下，如果启用了 llvm，则使用 llvm，否则使用 stackvm 解释器。
params (Optional[Dict[str, list]]) – 将要绑定的输入 IRModule 的参数。
relax_pipeline (str = "default") – 要使用的 Relax 编译 pipeline。
tir_pipelinie (str = "default") – 要使用的 TIR 编译 pipeline。
exec_mode ({"bytecode", "compiled"}) – 执行模式。
system_lib (Optional[bool]) – 是否构建系统库，该库被静态打包并自动将生成的函数注册到系统。默认情况下，根据目标自动检测。

返回:

ex – 可以被虚拟机加载的可执行文件。

返回类型:

tvm.relax.Executable

示例

class InputModule:
    @R.function
    def foo(x: Tensor((3, 4), "float32"), y: Tensor((3, 4), "float32")):
        z = R.add(x, y)
        return z

mod = InputModule
target = tvm.target.Target("llvm", host="llvm")
ex = tvm.compile(mod, target)

class tvm.relax.VMExecutable(mod: Module)

由 VM 编译器或 ExecBuilder 发出的虚拟机可执行对象。

stats() → str: 打印可执行文件的详细统计信息。

as_text() → str: 以文本格式打印指令。

as_python() → str: 以 python 程序格式打印指令。

class tvm.relax.DataflowBlockRewrite(dfb: DataflowBlock, root_fn: Function)

绑定/语句级数据流块重写器。

注释

由于 TVM AST 节点的不可变和写入时复制特性，重写不是就地完成的。相反，会创建一个新的 DataflowBlock 并通过 mutated_dfb 返回。类似地，它的新根 Function 会被创建并通过 mutated_root_fn 返回。要为 IRModule 应用此更改，请使用 mutate_irmodule，它会重写构造函数中注册的旧函数。

replace_all_uses(old_var: Var, new_var: Var) → None

将 old_var 的所有用法替换为 new_var。

参数:

old_var (relax.Var) – 要替换的旧变量。
new_var (relax.Var) – 用于替换的新变量。

add(expr: RelaxExpr, name: str | None = None, is_dfvar: bool = False) → None

向 DataflowBlock 添加一个新语句，并自动生成变量名。

参数:

expr (Expr) – 要添加的表达式。
name (Optional[str], optional) – 变量名，默认为 None
is_dfvar (bool, optional) – 变量类型，默认为 False

注释

如果未给出变量名，则将自动生成一个 “tmp${COUNTER}” 形式的变量名。如果 is_dfvar 为 True，则变量类型将为 DataflowVar，否则为 Var。作为 relax.Var 意味着变量是 DataflowBlock 的输出变量。而作为 DataflowVar 意味着变量是 DataflowBlock 的内部变量。

remove_unused(var: Var, allow_undef=False) → None

仅当语句的变量定义未使用时，才通过其变量定义删除语句。

参数:

var (relax.Var) – 未使用的变量定义。
allow_undef (bool, optional) – 是否允许 var 是未定义的变量，默认为 False

Raises:

TVMError if the variable is used or undefined (allow_undef=False). – 如果变量被使用或未定义（allow_undef=False），则抛出 TVMError。

remove_all_unused() → None

删除所有未使用的变量。

注释

这也可以删除其他 DataflowBlock 中未使用的变量。

mutated_dfb() → DataflowBlock: 返回已变异的 DataflowBlock。

mutated_root_fn() → Function: 返回已变异的根函数。

mutate_irmodule(irmodule: IRModule) → IRModule

返回一个更新后的 IRModule，方法是将旧函数替换为已变异的根函数。

参数:: irmodule (tvm.IRModule) – 要更新的基础 IRModule。
返回:: 更新后的 IRModule。
返回类型:: tvm.IRModule