tvm.relax.transform

Relax 转换。

tvm.relax.transform.AdjustMatmulOrder()

将 x*(A*B) 重排序为 (x*A)*B

用于优化 LoRA 计算，其中 matmul(x, LoraA*LoraB) 可以计算为 matmul(matmul(x, LoraA), LoraB)，从而减少总内存使用量。

返回:

ret – 相应的 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.AllocateWorkspace() → Pass

分配工作空间，用一个足够大的张量表示，以容纳所有需要临时存储的外部函数，并将其附加到外部函数的参数中。

外部函数可以通过 kWorkspaceSize 属性指定其工作空间需求。

返回:

ret – 用于分配工作空间的已注册 pass。

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.AlterOpImpl(op_impl_map: Dict[str, PrimFunc], op_buffer_transforms: Dict[str, List[IndexMap | Callable]], op_buffer_axis_separators: Dict[str, List[axis_separator | Callable]], op_buffer_input_axis_separators: Dict[str, List[axis_separator | Callable]])

替换所有具有匹配 ‘operator_name’ 属性的 PrimFunc，使用可能在 i/o 缓冲区上具有不同布局的替换 PrimFunc。i/o 缓冲区上的布局转换存在于 op_buffer_transforms 映射中。在被替换的 PrimFunc 的调用站点中插入布局转换，以将 i/o 张量转换为新 PrimFunc 期望的布局。

参数:

• op_impl_map (Dict[str, PrimFunc]) – op_kind 到 PrimFunc 的映射

• op_buffer_transforms (Dict[str, List[Union[IndexMap, Callable]]) – op_kind 到每个缓冲区的布局转换映射

• op_buffer_axis_separators (Dict[str, List[Union[IndexMap.AXIS_SEPARATOR, Callable]]]) – op_kind 到每个 index_map 的 axis_separator

• op_buffer_input_axis_separators (Dict[str, List[Union[IndexMap.AXIS_SEPARATOR, Callable]]]) – op_kind 到输入 index_map 的 axis_separator

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.AnnotateTIROpPattern() → Pass

为 TIR 函数注解 Op Pattern Kind

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.AttachAttrLayoutFreeBuffers() → Pass

将布局自由缓冲区附加到 tir::PrimFunc。

此 pass 用于根据 relax 函数中的函数使用情况，将布局自由缓冲区附加到 tir::PrimFunc。目前，布局自由缓冲区是模型权重和 relax 常量。

请注意，我们建议在此 pass 之前应用 CanonicalizeBindings。

返回:

ret – 用于附加布局自由缓冲区的已注册 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.AttachGlobalSymbol() → Pass

将 global_symbol 附加到 Relax 函数和 TIR PrimFunc 以进行代码生成。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.BindParams(func_name: str, params: Dict[str | Var, NDArray | ndarray]) → Pass

将模块的函数的参数绑定到常量张量。

参数:

• func_name (str) – 要绑定的函数名称

• params (Dict[Union[str,relax.Var], Union[tvm.runtime.NDArray, np.ndarray]]) – 从参数或参数名称到常量张量的映射。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.BindSymbolicVars(binding_map: Mapping[str | Var, PrimExpr], func_name: str | None = None) → Pass

将模块的函数的参数绑定到常量张量。

参数:

• binding_map (Mapping[Union[str, tvm.tir.Var], tvm.tir.PrimExpr]) – 从符号变量名到整数的映射。

• func_name (Optional[str]) – 要绑定的函数名称。如果为 None（默认值），则将更新模块中的所有函数。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.BundleModelParams(param_tuple_name: str | None = None) → Pass

将多个模型参数捆绑到单个元组参数中

对于每个函数，如果该函数具有 “num_input” 属性，则在运行时参数和编译时权重之间进行分隔。运行时参数（例如，激活）是前 num_input 个参数，其余的是编译时权重。

参数:

param_tuple_name (Optional[str]) – 元组参数的名称。如果未指定，则默认为 “model_params”。

返回:

ret – 用于捆绑模型参数的已注册 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.CallTIRRewrite() → Pass

为 call_tir 和 call_dps_packed 执行显式张量分配。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.CanonicalizeBindings() → Pass

规范化变量定义（例如，如果存在 y = x 和 z = y，则用 x 替换 y 和 z 的用法）。还简化匹配 cast 节点（消除冗余检查）和元组索引。

最好与常量折叠和消除未使用的定义结合使用。

注意：如果数据流变量仅在绑定到数据流块输出变量（即，非数据流变量）时使用，则此 pass 也会删除数据流变量，并将输出变量的绑定替换为数据流变量的直接定义。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.CombineParallelMatmul(check=None)

将共享相同 LHS 矩阵的多个 matmul 运算符合并为一个，然后进行切片。当树中的所有 matmul 分支都具有相同的融合操作集时，融合操作将应用于合并后的 matmul 输出，然后再进行切片。

目前，仅支持有限的融合操作集。它包括 bias add、relu、gelu、gelu_tanh 和 silu 激活。

参数:

check (Callable[[relax.Var, List[relax.Var], List[relax.Var], Dict[relax.Var, Expr]], bool]) – 用于过滤掉不需要的分支的函数，签名为 (input, [rhs], [bias], binding) -> bool。

返回:

ret – 相应的 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.ComputePrimValue() → Pass

计算所有 R.prim_value 实例

虽然高级 relax 可以包含关于其符号变量的表达式，但这些表达式不能在 relax 中本地计算。为了为符号表达式（例如 R.prim_value(N*N)，其中 N 是符号变量）提供值，此 pass 生成一个 PrimFunc，可以在其中计算表达式。然后更新 relax 图以包含对该 PrimFunc 的调用，以代替原始的 R.prim_value(expr)。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.ConvertLayout(desired_layouts: Dict[str, List[str]]) → Pass

自动布局转换 pass。

参数:

desired_layouts (Dict[str, List[str]]) – conv2d 运算符的期望布局是从运算符名称到期望的特征图、权重和输出的期望布局的映射。例如，如果我们想将 conv2d 的布局从 NCHW 转换为 NHWC，我们可以将 conv2d 的期望布局设置为 {"relax.nn.conv2d": ["NHWC", "OHWI"]}。

返回:

ret – 用于布局转换的已注册 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.ConvertToDataflow(min_size: int = 2) → Pass

一个 pass，用于将绑定块内的连续数据流操作转换为数据流块。

注意：ConvertToDataflow 可能需要先调用。

参数:

min_size (int) – pass 需要提取新块的连续数据流绑定的最小数量。

返回:

ret – 该 pass。

返回类型:

tvm.ir.transform.Pass

class tvm.relax.transform.DataflowBlockPass

一个 pass，用于处理模块中的每个 tvm.relax.DataflowBlock。

tvm.relax.transform.DataflowUseInplaceCalls() → Pass

Pass，用于将可以就地完成的运算符的调用（通常，这些是元素级操作）更改为就地实现。支持的运算符将被替换为 call_tir_inplace 的调用，该调用调用这些运算符的就地 PrimFunc 实现（这些实现基于这些运算符的合法化）。

注意：可能需要先调用 ConvertToDataflow 以提供数据流块。

返回:

ret – 该 pass

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.DeadCodeElimination(entry_functions: List[str] | None = None) → Pass

删除 IRModule 中的死代码。目前它删除

1. 未使用的本地 VarBinding（其中绑定的变量未使用且未使用任何不纯操作）。

2. 模块中未使用的 Relax 函数。我们检测从入口函数开始的调用链，并删除所有未使用的函数。

任何留空的绑定块都将由 normalizer 删除。

注释

对于函数式 DCE，请使用 py:func:tvm.relax.analysis.remove_all_unused。

参数:

entry_functions (Optional[List[str]]) – 要从其开始的入口函数集。

返回:

ret – 已注册的 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.DecomposeOpsForInference(func_name: str | None = None) → Pass

分解由推理期间的其他运算符组成的复合运算符。例如，批归一化的结果（三元组）将被简化。Attention、tensor_to_shape 等也可以分解为许多简化的运算符。

参数:

func_name (Optional[str]) – 指定函数的名称。如果未指定，则 pass 将在所有函数中运行。

返回:

ret – 已注册的 pass

返回类型:

tvm.transform.Pass

tvm.relax.transform.DecomposeOpsForTraining(func_name: str | None = None) → Pass

分解由训练期间的其他运算符组成的复合运算符。例如，批归一化的结果（三元组）将被简化。Attention、tensor_to_shape 等也可以分解为许多简化的运算符。

参数:

func_name (Optional[str]) – 指定函数的名称。如果未指定，则 pass 将在所有函数中运行。

返回:

ret – 已注册的 pass

返回类型:

tvm.transform.Pass

tvm.relax.transform.EliminateCommonSubexpr(call_only=False) → FunctionPass

消除函数内的公共子表达式。

注意：对于嵌套函数，此 pass 在这些函数内部执行 CSE

参数:

call_only (bool) – 如果为 True，则启用仅消除调用节点。

返回:

ret – 用于消除公共子表达式的已注册 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.ExpandMatmulOfSum()

将 matmul(x, A+B) 展开为 matmul(x,A) + matmul(x,B)

如果任一操作数可以在编译时完全计算（仅取决于 kNumInput 之后的函数参数），则会抑制此展开。

用于优化 LoRA 计算，其中 matmul(x, Base + LoraA*LoraB) 可以展开为 matmul(x, Base) + matmul(x, LoraA*LoraB)，从而可以使用 CombineParallelMatmul 进行优化。

返回:

ret – 相应的 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.ExpandTupleArguments() → Pass

将元组参数展开为内部函数

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.FewShotTuning(valid_count: int = 1, benchmark: bool = False) → Pass

此 Pass 专为静态形状 PrimFunc 的小样本调优而设计。它检查 PrimFunc 中的所有块，并基于 MetaSchedule 调度规则执行循环融合、分割和其他转换，但直接从搜索空间中采样，而不是使用调优算法。用户可以指定要尝试的有效计数数量以及是否使用 runner 进行基准测试。

参数:

• valid_count (int) – 要尝试的有效计数数量。

• benchmark (bool) – 是否使用 runner 进行基准测试。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.FoldConstant() → Pass

折叠数据流块内的常量表达式。

注意：可能需要先调用 ConvertToDataflow 以提供数据流块。

返回:

ret

返回类型:

tvm.ir.transform.Pass

class tvm.relax.transform.FunctionPass

一个 Pass，作用于模块中的每个 tvm.relax.Function。函数 Pass 类应通过 function_pass 创建。

tvm.relax.transform.FuseOps(fuse_opt_level=-1) → Pass

此 Pass 根据 Pass 实现中描述的融合算法，对 Relax 函数的数据流块中的绑定进行分组，并为每个组生成一个新的分组 Relax 函数。通过将绑定分组到新的 Relax 函数中，我们将被操作函数中的绑定替换为对新分组函数的函数调用。

名为 “FuseTIR” 的后续 Pass 将为每个分组函数生成一个 TIR PrimFunc。

注意：可能需要先调用 ConvertToDataflow 以提供数据流块。

参数:

fuse_opt_level (int) – 融合优化级别。-1 表示级别将从 Pass 上下文中推断。

返回:

ret – 运算符融合的已注册 Pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.FuseOpsByPattern(patterns: List[FusionPattern | Tuple], bind_constants: bool = True, annotate_codegen: bool = False, entry_functions: List[str] | None = None) → Pass

将模式匹配应用于给定模块中的每个函数，并将匹配的表达式分组到一个新函数中。

最终结果类似于 FuseOps，但融合完全由提供的模式驱动。

注意：仅在数据流块内操作。可能需要先调用 ConvertToDataflow。

参数:

• patterns (List[Union[FusionPattern, Tuple]]) –

  要匹配的模式列表。模式的顺序决定了它们被匹配的优先级顺序。优先级较高的模式应在列表中靠前。

  除了 FusionPattern，还可以将元组作为此列表的项传递。模式将通过 FusionPattern(*item) 构建

• bind_constants (bool) – 是否将绑定的常量保留在分组函数中。

• annotate_codegen (bool) –

  如果为 True，则用另一个函数包装每个创建的复合函数，该函数的主体仅包含对复合函数的调用，并使用 “Codegen” 和 “global_symbol” 属性注释外部函数。“Codegen” 属性设置为相应模式名称的前缀。例如，如果模式名称为 “dnnl.conv2d_relu”，则为 “dnnl”。

  如果要将创建的复合函数卸载到外部后端，而不使用 MergeCompositeFunctions Pass，则此项必须为 True。

• entry_functions (Optional[List[str]]) – 要从其开始的入口函数集。

返回:

ret – 基于模式融合的已注册 Pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.FuseTIR() → Pass

如果可能，将原始 relax 函数融合到更大的 TIR 函数中

返回:

ret – tir 融合的已注册 Pass。

返回类型:

tvm.transform.Pass

class tvm.relax.transform.FusionPattern(name: str, pattern: DFPattern, annotation_patterns: Mapping[str, DFPattern] | None = None, check: Callable[[PatternCheckContext], bool] | None = None, attrs_getter: Callable[[Dict[str, RelaxExpr]], Dict[str, str]] | None = None)

FuseOpsByPattern 使用的模式。它主要是 DFPattern，但带有其他信息以在融合 Pass 期间提供帮助。

参数:

• name (str) – 模式的名称。通常它以后端名称开头，例如 ‘cutlass.matmul’。

• pattern (DFPattern) – 数据流模式，将用于匹配可由外部后端处理的表达式。

• annotation_patterns (Mapping[str, DFPattern]) – 用于从模式匹配结果中提取重要表达式的映射。此映射中的所有 DFPattern 都应是 pattern 的一部分。

• check (Callable[[PatternCheckContext], bool]) – 用于检查匹配结果是否被接受的函数。

tvm.relax.transform.Gradient(func_name: str, require_grads: Var | List[Var] | None = None, target_index: int = 0) → Pass

反向模式自动微分。

此 Pass 将微分 IRModule 中的一个函数。现在，输入函数必须只有一个数据流块（可能需要先调用 ConvertToDataflow）。

对于由 func_name 指定的给定函数，它会生成一个新函数，名称为 func_name + “_adjoint”。新函数计算 **微分目标** 相对于原始函数的 require_grads 指定的参数的梯度。

如果函数只有一个返回值，则返回值将指定为目标。如果函数有多个返回值，则目标将被指定为第 target_index 个返回值。目标必须是标量（0 维张量）。

新函数将类似于

@R.function
def main_adjoint(original_parameters):
    with R.dataflow():
        # the bindings of the original function
        ...
        # calculating the gradients
        ...
        R.output(original_outputs, grad_1, grad_2, ...)
    return (original_return_value, (grad_1, grad_2, ...))

此 AD Pass 还支持检查点，如 “Training deep nets with sublinear memory cost.” - Chen, Tianqi, et al. (2016) 中所述。有关更多详细信息，请参阅 tvm.relax.testing.nn.checkpoint。

参数:

• func_name (str) – 特定函数的名称。

• require_grads (Optional[Union[relax.Var, List[relax.Var]]]) – 需要其伴随的 relax 变量。必须是给定函数的参数，并且不应重复。如果未指定，则将计算所有参数的伴随。

• target_index (int) – 如果指定的函数有多个返回值，请指定返回值索引作为目标。如果未指定，则第一个返回值将作为目标。

返回:

ret – Pass。

返回类型:

tvm.ir.transform.Pass

示例

以下代码显示了如何使用此 Pass

@I.ir_module
class Module:
    @R.function
    def main(
        x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
    ) -> R.Tensor((), dtype="float32"):
        with R.dataflow():
            lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
            # use R.sum to reduce the tensor to a scalar
            lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
            R.output(lv2)
        return lv2

After = relax.transform.Gradient("main")(Module)

Gradient Pass 之后的模块将是

@I.ir_module
class After:
    @R.function
    def main(
        x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
    ) -> R.Tensor((), dtype="float32"):
        with R.dataflow():
            lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
            lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
            R.output(lv2)
        return lv2

    @R.function
    def main_adjoint(
        x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
    ) -> R.Tuple(
        R.Tensor((), dtype="float32"),
        R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")),
    ):
        with R.dataflow():
            # original bindings
            lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
            lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
            # bindings w.r.t. intermediate variables
            lv2_adjoint: R.Tensor((), dtype="float32") = R.ones((), dtype="float32")
            lv1_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(
                lv2_adjoint, (3, 3)
            )
            # bindings w.r.t. parameters
            x_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
            y_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
            R.output(lv2, x_adjoint, y_adjoint)
        # return value: (orig_return_values, tuple(adjoints))
        return (lv2, (x_adjoint, y_adjoint))

第二个示例返回多个值，并使用 target_index 指定目标

@I.ir_module
class Module:
    @R.function
    def main(
        x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
    ) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tensor((), dtype="float32")):
        with R.dataflow():
            lv1: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
            lv2: R.Tensor((), dtype="float32") = R.sum(y, axis=None, keepdims=False)
            R.output(lv1, lv2)
        return (lv1, lv2)

After = relax.transform.Gradient("main", target_index=1)(Module)

Gradient Pass 之后的模块将是

@I.ir_module
class Module:
    @R.function
    def main(
        x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
    ) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tensor((), dtype="float32")):
        with R.dataflow():
            lv1: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
            lv2: R.Tensor((), dtype="float32") = R.sum(y, axis=None, keepdims=False)
            R.output(lv1, lv2)
        return (lv1, lv2)

    @R.function
    def main_adjoint(
        x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
    ) -> R.Tuple(
        R.Tuple(R.Tensor((), dtype="float32"), R.Tensor((), dtype="float32")),
        R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")),
    ):
        with R.dataflow():
            # original bindings
            lv1: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
            lv2: R.Tensor((), dtype="float32") = R.sum(y, axis=None, keepdims=False)
            # bindings w.r.t. intermediate variables
            # gradient of intermediate variables that is not related to the target will not
            # be calculated
            lv2_adjoint: R.Tensor((), dtype="float32") = R.ones((), dtype="float32")
            # bindings w.r.t. parameters
            x_adjoint: R.Tensor((3, 3), dtype="float32") = R.zeros((3, 3), dtype="float32")
            y_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(
                lv2_adjoint, (3, 3)
            )
            R.output(lv1, lv2, x_adjoint, y_adjoint)
        # return value: (orig_return_values, tuple(adjoints))
        return ((lv1, lv2), (x_adjoint, y_adjoint))

tvm.relax.transform.InlinePrivateFunctions() → Pass

内联所有私有的 relax 函数

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.KillAfterLastUse() → Pass

在最后一次使用后删除所有张量/存储对象

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.LambdaLift() → Pass

一个将局部函数提升为全局函数的 Pass。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.LazyGetInput() → Pass

一个延迟请求输入的 Pass。

在许多情况下，模型权重的尺寸超过了 GPU 上的可用内存。在这些情况下，接受所有模型权重作为参数的函数将无法被调用。在这些情况下，必须在函数需要时加载参数，并在不再需要时卸载参数。

此 Pass 改变函数，使得所有模型权重（第一个 func.attrs[“num_input”] 参数之后的参数）按需加载。该函数不再接受权重作为函数参数，而是接受一个回调参数，该回调参数可以根据需要加载每个参数。回调接受两个参数，第一个是模型权重的索引，第二个是参数的名称。回调应返回指定的参数。

@R.function
def before(A: R.Tensor([16,32],"float32")):
    ...

@R.function
def after(fget_param: R.Callable([R.Prim('int64'), R.Object], R.Object)):
    A_untyped = fget_param(0, R.str('A'))
    A = R.match_cast(A_untyped, R.Tensor([16,32], "float32")
    ...

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.LazySetOutput() → Pass

一个在输出可用时设置函数输出的 Pass

在许多情况下，模型权重的尺寸超过了 GPU 上的可用内存。在这些情况下，将所有模型权重作为单个返回值生成的函数将无法被调用。在这些情况下，参数必须在生成时返回，从 GPU 卸载（或保存到磁盘），然后再生成其他输出。

此 Pass 改变函数，使得函数的所有输出在可用时返回。该函数接受一个额外的回调参数，该回调参数在函数的每个输出处被调用。回调接受两个参数，第一个是生成的输出元组的索引（如果输出不是元组，则为零），第二个是值本身。

@R.function
def before(args):
    ...
    return (A, B)

@R.function
def after(args, fset_param: R.Callable([R.Prim('int64'), R.Object])):
    ...
    fset_param(0, A)
    ...
    fset_param(1, B)
    ...
    return ()

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.LegalizeOps(customize_legalize_map: Dict[str, Callable[[BlockBuilder, Call], RelaxExpr]] | None = None, enable_warning: bool = False)

将 Relax 函数中的高级运算符调用合法化为带有相应低级 TIR PrimFunc 的 call_tir。

对于每个高级运算符，我们将合法化它的方式注册为一个函数，该函数接受上下文 BlockBuilder 和要合法化的 relax.Call 作为输入，并返回合法化的调用。这里的输入 BlockBuilder 主要用于将 call_te 创建的 PrimFunc 添加到上下文 IRModule 中。

每个运算符的合法化函数都注册为运算符的属性（属性键为 FLegalize）。

此 Pass 为用户提供了自定义性，可以使用他们自己的运算符合法化函数。该 Pass 采用可选的自定义映射，键为运算符名称 (str)，值为函数 (LegalizeFunc)。默认的合法化函数将被自定义函数覆盖。

参数:

• customize_legalize_map (Optional[Dict[str, LegalizeFunc]]) – 自定义运算符合法化函数映射。自定义函数将覆盖默认函数。

• enable_warning (bool) – 一个布尔值，指示是否为运算符的合法化函数未注册的 CallNode 打印警告。默认情况下，我们不打印警告。

返回:

ret – 已注册的 pass

返回类型:

tvm.transform.Pass

示例

以下代码显示了如何使用此 Pass

# Define the pass input IRModule
@tvm.script.ir_module
class Module:
    @R.function
    def main(
        x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")
    ) -> R.Tensor((2, 3), "float32"):
        z: R.Tensor((2, 3), "float32") = R.add(x, y)
        r: R.Tensor((2, 3), "float32") = R.multiply(y, z)
        return r

# Define the customized legalization function for "relax.add"
def customize_legalize_add(bb: relax.BlockBuilder, call: relax.Call) -> relax.Expr:
    from tvm import topi
    return bb.call_te(topi.add, call.args[1], call.args[0])

# Apply the pass with the customized function to the module.
mod = LegalizeOps({"relax.add": customize_legalize_add})(Module)

通过 mod.show() 打印输出结果，我们可以看到合法化后的 IRModule 变为

@tvm.script.ir_module
class Module:
    @R.function
    def main(
        x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")
    ) -> R.Tensor((2, 3), "float32"):
        z = R.call_tir(add, (y, x), (2, 3), dtype="float32")
        r = R.call_tir(multiply, (y, z), (2, 3), dtype="float32")
        return r

    @T.prim_func
    def add(
        A: T.Buffer((2, 3), "float32"),
        B: T.Buffer((2, 3), "float32"),
        T_add: T.Buffer((2, 3), "float32"),
    ):
        T.func_attr({"tir.noalias": True})
        for ax0, ax1 in T.grid(2, 3):
            with T.block("T_add"):
                v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])
                T.writes(T_add[v_ax0, v_ax1])
                T_add[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[v_ax0, v_ax1]

    @T.prim_func
    def multiply(
        A: T.Buffer((2, 3), "float32"),
        B: T.Buffer((2, 3), "float32"),
        T_multiply: T.Buffer((2, 3), "float32"),
    ):
        T.func_attr({"tir.noalias": True})
        for ax0, ax1 in T.grid(2, 3):
            with T.block("T_multiply"):
                v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])
                T.writes(T_multiply[v_ax0, v_ax1])
                T_multiply[v_ax0, v_ax1] = A[v_ax0, v_ax1] * B[v_ax0, v_ax1]

tvm.relax.transform.LiftTransformParams(shared_transform: bool | List[str] = False) → Pass

提升函数的参数的转换。

当函数的某些输入被标记为 “parameters”（模型权重）时，此 Pass 识别参数的转换，并将它们提升到一个名为 transform_params 的单独函数。transform_params 接受原始参数的元组作为输入，并返回转换后的参数的元组。原始函数将被重写以接受转换后的参数的元组作为输入。

用户应在运行时调用 transform_params 函数，并将转换后的参数作为输入传递给原始函数。

参数:

shared_transform (Union[bool, List[str]]) –

指示将如何生成参数转换函数

• False （默认）：将为每个具有 “num_input” 属性的函数生成单独的参数转换函数。

• True：将生成单个参数转换函数，其中包含所有具有 “num_input” 属性的函数共有的预处理步骤。

• List[str]：将生成单个参数转换函数，其中包含名称在列表中的每个函数共有的预处理步骤。传递所有具有 “num_input” 属性的函数的列表或空列表等同于传递 True。

返回:

ret – 用于提升参数转换的已注册 Pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.LowerAllocTensor() → Pass

降低 R.builtin.alloc_tensor 的剩余实例

静态内存规划器删除 R.builtin.alloc_tensor 的静态实例，并替换为 R.memory.alloc_storage 和 R.memory.alloc_tensor。但是，R.builtin.alloc_tensor 仍然保留用于任何动态分配。

此转换将任何剩余的 R.builtin.alloc_tensor 实例替换为 R.memory.alloc_storage 和 R.memory.alloc_tensor。如果不存在 R.builtin.alloc_tensor，则此 Pass 无效。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.LowerRuntimeBuiltin() → Pass

将通用内联函数降低为 VM 内联函数。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.MergeCompositeFunctions() → Pass

将 FuseOpsByPattern 创建的一个或多个复合函数分组到一个新函数中。新函数将使用 “Codegen” 和 “global_symbol” 属性进行注释，并且旨在卸载到外部后端。

返回:

ret – 用于合并复合函数的已注册 Pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.MetaScheduleApplyDatabase(work_dir: str | None = None, enable_warning: bool = False) → Pass

应用来自调优数据库的最佳调度。

参数:

• work_dir (Optional[str]) – 工作目录，用于在未提供数据库时推断默认数据库（当用户传递数据库时将被忽略）

• enable_warning (bool) – 一个布尔值，指示是否为数据库中未显示的 TIR 函数打印警告。默认情况下，我们不打印警告。

返回:

ret – 已注册的 pass

返回类型:

tvm.transform.Pass

tvm.relax.transform.MetaScheduleTuneIRMod(params: Dict[str, NDArray], work_dir: str, max_trials_global: int, max_trials_per_task: int | None = None, op_names: List[str] | None = None) → Pass

使用 MetaSchedule 调优 Relax IRModule。

参数:

• params (Dict[str, NDArray]) – 模型参数

• work_dir (str) – 工作目录

• max_trials_gloabl (int) – 调优允许的最大总试验次数

• max_trials_per_task (int) – 每个任务的最大试验次数

• op_names (Optional[List[str]]) – 用于指定要调优的算子的算子名称列表。当为 None 时，所有算子都将被调优。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.MetaScheduleTuneTIR(work_dir: str, max_trials_global: int) → Pass

使用 MetaSchedule 调优 TIR。 :param work_dir: 工作目录 :type work_dir: str :param max_trials_gloabl: 调优允许的最大总试验次数 :type max_trials_gloabl: int

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.Normalize() → Pass

将 Relax IR 转换为 normal form，即表达式被规范化（没有嵌套，因此 AST 处于 ANF 形式），并且所有表达式的 checked_type_ 和 shape_ 均可用。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.NormalizeGlobalVar() → Pass

可能重命名 IRModule 中的 GlobalVar 以确保以下属性

1. (不变性) 首先确保每个公共函数都具有与其 “global_symbol” 属性相同的名称 2. 为了确保 1.，我们可能需要重命名具有冲突名称的私有函数；3. 最后，每个 GlobalVar 的名称在 IRModule 中都是唯一的。

返回:

ret

返回类型:

tvm.ir.transform.Pass

class tvm.relax.transform.PatternCheckContext

FusionPattern.check 检查函数的输入。

参数:

• matched_expr (Expr) – 与 FusionPattern.pattern 匹配的表达式。

• annotated_expr (Mapping[str, Expr]) – 一个映射，其中包含 FusionPattern.annotation_patterns 中子模式匹配的所有表达式。

• matched_bindings (Mapping[relax.Var, Expr]) – 从变量到其值的映射。它包含来自 FuseOpsByPattern 融合的绑定的变量。

• var_usages (Mapping[relax.Var, Sequence[relax.Var]]) – 一个映射，将变量定义映射到一组用途。它包含函数中使用的所有变量。

• value_to_bound_var (Mapping[Expr, relax.Var]) – 从值到其绑定变量的映射。它不包含匹配表达式之后的变量。

tvm.relax.transform.RealizeVDevice() → Pass

传播虚拟设备信息。

返回:

ret – 已注册的 pass

返回类型:

tvm.transform.Pass

tvm.relax.transform.RemovePurityChecking() → Pass

在模块中的所有纯函数上激活 relax.force_pure，并将所有纯覆盖操作解包为正常版本。

这实际上意味着将不再进行纯度跟踪，这对于底层代码生成很有用。

返回:

ret – Pass。

返回类型:

tvm.ir.transform.Pass

注意

应在 ToNonDataflow() 之后使用

tvm.relax.transform.RemoveUnusedOutputs() → Pass

从内部函数中删除未使用的输出

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.RemoveUnusedParameters() → Pass

删除内部函数中未使用的参数

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.ReorderPermuteDimsAfterConcat()

将 concat(permute_dims(A), permute_dims(B)) 重排序为 permute_dims(concat(A,B))

在 CombineParallelMatmul 之后优化计算很有用。优化的 nn.Linear 实现的模式查找 matmul(activations, permute_dims(weights))。在 CombineParallelMatmul 之后，matmul(activations, concat(permute_dims(A), permute_dims(B))) 不再匹配此模式。重新排列成 matmul(activations, permute_dims(concat(A,B))) 恢复了模式匹配。

返回:

ret – 相应的 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.ReorderTakeAfterMatmul()

将 matmul(x, take(weights, indices)) 重排序为 take(matmul(x,weights),indices)

对于优化 LoRA 计算很有用，其中可以将多个 LoRA 批量处理在一起。

返回:

ret – 相应的 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.RewriteCUDAGraph() → Pass

重写 Relax 模块以使用 CUDA 图执行。此 pass 识别可以使用 CUDA 图执行的区域，并将它们提升为新的函数以进行运行时图捕获。

返回:

ret – 用于重写 cuda 图的已注册 pass

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.RewriteDataflowReshape() → Pass

将所有类 reshape 的 call_tir 转换为 VM reshape 运算符调用。VM reshape 运算符调用将在运行时进一步降级为 CreateView 操作，而不是执行实际的数据复制。此处 “类 reshape” 包括 reshape、expand_dims、flatten 等。

注意：仅在 dataflow 块中操作。可能需要首先调用 ConvertToDataflow。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.RunCodegen(target_options: dict | None = None, entry_functions: List[str] | None = None) → Pass

生成带有注解代码生成和全局符号的 runtime::Module。

参数:

• target_options (Optional[dict]) – 目标名称和编译选项对

• entry_functions (Optional[List[str]]) – 要从其开始的入口函数集。

返回:

ret – 用于删除未使用函数的已注册 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.SplitCallTIRByPattern(patterns: List[PrimFunc], fcodegen: Callable) → Pass

将 PrimFunc 拆分为 2 个部分：第一部分是 TIR PrimFunc，它

与某些模式匹配，第二部分是原始 PrimFunc 的其余部分。它将调用 fcodegen 为匹配的模式生成代码，以将其替换为 ExternFunc 调用。

参数:

• patterns (List[PrimFunc]) – 要匹配的模式列表。

• fcodegen (Callable[[List[MatchResult]], List[Object]]) – 用于为匹配模式生成代码的函数。

返回:

ret – 用于拆分 call_tir 的已注册 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.SplitLayoutRewritePreproc() → Pass

将 TIR 布局重写拆分为多个 TIR 函数。此 pass 在 meta_schedule 调优后用于预打包权重。

返回:

ret – 用于拆分 TIR 布局重写的已注册 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.StaticPlanBlockMemory() → Pass

BindingBlock 级别的静态内存规划 pass。该 pass 将尽最大努力重用已分配的内存，以减少已分配内存的总大小。

该 pass “支持” 动态形状，通过 TIR 变量上限注解的方式。我们可以选择性地将属性 “tir_var_upper_bound” 注解到 Relax 函数。属性值是从字符串到整数的字典，表示 TIR 变量的名称到 TIR 变量的上限值。注意：注解的上限属性仅适用于函数签名中的 TIR 变量，以提高清晰度。

例如，我们可以使用 R.func_attr({"tir_var_upper_bound": {"n": 1024}}) 注解 Relax 函数。这意味着函数签名中名为 “n” 的变量的最大值将具有上限 1024。我们将在内存规划期间使用 1024 作为其值。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.ToMixedPrecision(out_dtype='float32', fp16_input_names: List[str] | None = None) → Pass

自动混合精度 pass。目前，该 pass 假设输入模块仅为 fp32，并将自动将 fp32 转换为 fp16 以用于某些操作。

注意：主要在 dataflow 块内操作。可能需要首先调用 ConvertToDataflow。

参数:

• out_dtype (str) – gemm/conv 的输出数据类型，即累加器的数据类型。

• fp16_input_names (List[str]) – 应该变为 fp16 数据类型的函数参数的名称列表。函数签名将相应地更改。

返回:

ret – 用于混合精度的已注册 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.ToNonDataflow() → Pass

将所有 dataflow 结构转换为 non-dataflow 版本。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.TopologicalSort(order='depth-first', direction='from-inputs') → Pass

按照指定的顺序对 relax.Dataflow 块中的绑定进行排序

参数:

• order (str) – 应该发出绑定的顺序。允许的值为 “depth-first” 和 “breadth-first”。

• direciton (str) – 应执行排序的方向。允许的值为 “from-inputs” 和 “from-outputs”。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.UpdateParamStructInfo(sinfo_func: Callable[[Var], StructInfo | None])

更新参数的结构信息

更新参数的结构信息。内部绑定和函数返回类型将使用 relax 的结构推断规则进行更新。结构推断产生的错误将传播给用户。

参数:

sinfo_func (Callable[[relax.Var], Optional[StructInfo]]) – 一个函数，对于每个函数参数调用一次，并返回用于该参数的更新后的结构信息。如果函数返回 None，则不会修改该参数。

返回:

ret – 相应的 pass。

返回类型:

tvm.transform.Pass

tvm.relax.transform.UpdateVDevice(new_vdevice: VDevice, index: int) → Pass

更新虚拟设备。

参数:

• new_vdevice (tvm.ir.VDevice) – 新的虚拟设备。

• index (int) – 设备索引，指示将在哪个设备上执行更新。

返回:

ret – 修改虚拟设备的已注册 pass。

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.VMBuiltinLower() → Pass

将通用内联函数降低为 VM 内联函数。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.VMShapeLower(*, emit_err_ctx: bool = True) → Pass

降低符号形状和参数以及 match-cast 结构信息匹配。

参数:

emit_err_ctx (Optional[bool]) – 是否发出错误上下文字符串，为了测试目的可以关闭。

返回:

ret

返回类型:

tvm.ir.transform.Pass

tvm.relax.transform.dataflowblock_pass(pass_func=None, opt_level=None, name=None, required=None, traceable=False) → Callable | DataflowBlockPass

装饰一个 dataflowblock pass。

当提供 pass_func 时，此函数返回一个回调。否则，它返回使用给定优化函数创建的 dataflowblock pass。

参数:

• pass_func (Optional[Callable[(DataflowBlock, Module, PassContext) -> DataflowBlock]]) – 转换函数或类。

• opt_level (int) – 此 dataflowblock pass 的优化级别。

• name (Optional[str]) – dataflowblock pass 的名称。名称可以为空。在这种情况下，优化函数的名称将用作 pass 名称。

• required (Optional[List[str]]) – dataflowblock pass 依赖的 pass 列表。

• traceable (Boolean) – 布尔变量，指示 dataflowblock pass 是否可跟踪

返回:

create_dataflowblock_pass – 如果未提供 pass_func，将返回一个装饰器，否则返回装饰后的结果。返回的装饰器根据输入具有两种行为：当我们装饰一个 pass 函数时，将返回一个新的 DataflowBlockPass。当我们装饰一个类类型时，将返回一个新的 DataflowBlockPass 类。

返回类型:

Union[Callable, DataflowBlockPass]

示例

以下代码块装饰了一个 dataflowblock pass 类。

@relax.transform.dataflowblock_pass(opt_level=1)
class TestReplaceBinding:
    # Simple test function to replace the first VarBinding to another.

    def __init__(self):
        # create a new VarBinding
        m, n = tir.Var("m", "int64"), tir.Var("n", "int64")
        lv0 = relax.Var("lv1", relax.TensorStructInfo([m, n], "float32"))
        val = relax.const(np.random.rand(24, 56))
        self.new_binding = relax.VarBinding(lv0, val)

    def transform_dataflowblock(self, block, mod, ctx):
        # just for demo purposes
        # Replace the first binding in the DataflowBlock
        new_bindings = [self.new_binding, block.bindings[1]]
        new_block = relax.expr.DataflowBlock(new_bindings, block.span)
        return new_block

@tvm.script.ir_module
class InputMod:
    @R.function
    def f1(x: Tensor[(m, n), "float32"]):
        with relax.dataflow():
            lv0 = relax.multiply(x, x)
            gv0 = relax.add(x, x)
            relax.output(gv0)
        return gv0
# block_pass is now a special pass that replaces every
# first binding to the constant value binding
block_pass = TestReplaceBinding()
# now every first binding in DataflowBlock of InputMod
# is replaced by new_binding
updated_mod = block_pass(InputMod)

以下代码通过装饰用户定义的转换函数来创建 dataflowblock pass。

@relax.transform.dataflowblock_pass(opt_level=2)
def transform(block, mod, ctx):
    # my transformations here.
    return block

block_pass = transform
assert isinstance(block_pass, relax.transform.DataflowBlockPass)
assert block_pass.info.opt_level == 2

# Given a module m, the optimization could be invoked as the follwoing:
updated_mod = block_pass(m)
# Now transform should have been applied to every DataflowBlock in
# the provided module m. And the updated module will be returned.

tvm.relax.transform.function_pass(pass_func=None, opt_level=None, name=None, required=None, traceable=False) → Callable | FunctionPass

装饰一个函数 pass。

此函数在提供 pass_func 时返回一个回调。否则，它将返回使用给定优化函数创建的函数 pass。

参数:

• pass_func (可选[Callable[(Function, Module, PassContext) -> Function]]) – 转换函数或类。

• opt_level (int) – 此函数 pass 的优化级别。

• name (可选[str]) – 函数 pass 的名称。名称可以为空。在这种情况下，优化函数的名称将用作 pass 名称。

• required (可选[List[str]]) – 函数 pass 依赖的 pass 列表。

• traceable (Boolean) – 布尔变量，指示函数 pass 是否可追踪

返回:

create_function_pass – 如果未提供 pass_func，则将返回一个装饰器，否则返回装饰后的结果。返回的装饰器根据输入具有两种行为：当我们装饰一个 pass 函数时，将返回一个新的 FunctionPass。当我们装饰一个类类型时，将返回一个新的 FunctionPass 类。

返回类型:

Union[Callable, FunctionPass]

示例

以下代码块装饰了一个函数 pass 类。

@relax.transform.function_pass(opt_level=1)
class TestReplaceFunc:
    def __init__(self, new_func):
        self.new_func = new_func

    def transform_function(self, func, mod, ctx):
        # just for demo purposes
        # transform func to new_func
        return self.new_func

@R.function
def f1(x: Tensor[(m, n), "float32"]):
    return x

@tvm.script.ir_module
class InputMod:
    @R.function
    def f2(x: Tensor[(m, n), "float32"]):
        gv0 = relax.add(x, x)
        return gv0
# fpass is now a special pass that replaces every
# function to f1
fpass = TestReplaceFunc(f1)
# now every function in InputMod is replaced by f1
updated_mod = fpass(InputMod)

以下代码通过装饰用户定义的转换函数来创建一个函数 pass。

@relax.transform.function_pass(opt_level=2)
def transform(func, mod, ctx):
    # my transformations here.
    return func

function_pass = transform
assert isinstance(function_pass, relax.transform.FunctionPass)
assert function_pass.info.opt_level == 2

# Given a module m, the optimization could be invoked as the follwoing:
updated_mod = function_pass(m)
# Now transform should have been applied to every function in
# the provided module m. And the updated module will be returned.

class tvm.relax.transform.AttachExternModules(extern_modules: List[ExternModule])

将变量边界附加到每个 Relax 函数，这主要有助于内存规划。

class tvm.relax.transform.FastMathTransform(*args, **kwargs)

Pass，用于将昂贵的非线性函数转换为其快速但近似的对应函数。

class tvm.relax.transform.FuseTransposeMatmul(*args, **kwargs)

一个编译器 pass，用于融合转置 + 矩阵乘法。

class tvm.relax.transform.IPCAllReduceRewrite(allreduce_strategy: int)

将 all-reduce 操作重写为使用 IPC 内存的自定义 all-reduce 实现。

class tvm.relax.transform.LazyTransformParams(fget_item='get_item', fset_item='set_item', extra_get_item_params=None, extra_set_item_params=None)

将 transform_params 函数转换为惰性版本。（按需将输入加载到内存中，并在最后一次使用后立即释放它。）

注意：应在此 pass 之前调用 ToNonDataflow() 和 RemovePurityTracking()。

参数:

• fget_item (str) – get_item 函数的名称。

• fset_item (str) – set_item 函数的名称。

• extra_get_item_params (list of relax.Var) – get_item 函数的参数，索引除外。给定的参数将放在索引之前。例如，如果 extra_get_item_params 是 [param1, param2]，则 pass 将生成 call_packed(fget_item, [param1, param2, index])

• extra_set_item_params (list of relax.Var) – set_item 函数的参数，索引和值除外。给定的参数将放在索引和值之前。例如，如果 extra_set_item_params 是 [param1, param2]，则 pass 将生成 call_packed(fset_item, [param1, param2, index, value])

class tvm.relax.transform.LowerGPUIPCAllocStorage(*args, **kwargs)

降低 IPC 内存上的存储/张量分配。

class tvm.relax.transform.OptimizeLayoutTransform

Pass，用于删除 AlterOpImpl pass 引入的冗余转换布局运算符。

class tvm.relax.transform.FoldBatchnormToConv2D

将 BatchNorm 融合到其之前的 Conv2D。此优化是将 scale 折叠到 conv2d 权重中的 FoldScaleAxis 的特殊情况。当 FoldScaleAcis 增强以支持这种情况时，可以删除此 pass。

class tvm.relax.transform.RemoveRedundantReshape

转换 pass，用于删除冗余的 reshape 运算符