跳到主要内容
版本：0.21.0

tvm.tir.transform

所有 TIR 变换的命名空间

tvm.tir.transform.prim_func_pass(pass_func=None, opt_level:int|None= None, name:str|None= None, required:List[str] |None= None, traceable=False) → Callable | PrimFuncPass​

装饰一个函数传递。

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

  • 参数：
    • pass_func (可选*[*可调用[(tvm.tir.PrimFunc,IRModule,PassContext)–> tvm.tir.PrimFunc]] )：转换函数或类。
    • opt_level（int）：此模块传递的优化级别。
    • name*（可选*[ str]）：函数传递的名称。名称可以为空。在这种情况下，优化函数的名称将用作传递名称。
    • required（可选[List[str]]）：函数传递所依赖的传递列表。
  • 返回：create_function_pass：如果未提供 pass_func，则返回一个装饰器；否则返回装饰后的结果。返回的装饰器根据输入有两种行为：装饰一个 pass 函数时，将返回一个新的 FunctionPass。装饰一个类类型时，将返回一个新的 FunctionPass 类。
  • 返回类型： Union[Callable, FunctionPass]。

示例

下面的代码块装饰了一个函数传递类。

@tvm.tir.transform.prim_func_pass(opt_level=1)
class TestReplaceFunc:
    def **init**(self, new_func):
        self.new_func = new_func

    def transform_function(self, func, mod, ctx):
        # 仅用于演示
        # 将 func 转换为 new_func
        return self.new_func

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

@tvm.tir.transform.prim_func_pass(opt_level=2)
def transform(func, mod, ctx):
    # 我的转换操作写在这里。
    return func

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

# 给定一个模块 m，可以如下调用优化：
updated_mod = function_pass(m)
# 现在常量折叠（constant folding）应该已经应用到提供的模块 m 中的每个函数。
# 并且返回更新后的模块。

class tvm.tir.transform.PrimFuncPass​

tvm.tir.PrimFunc()对模块中的每个模块进行操作的传递。应通过 py:func: tvm.tir.transform.function_pass 创建函数传递类。

tvm.tir.transform.AnnotateDeviceRegions()​

注释应在设备上运行的位置。

插入 AttrStmt 节点，指定 PrimFunc 中应执行区域的目标。仅修改具有 tvm::attr::kTarget 属性且该目标定义了主机的函数。

tvm.tir.transform.AnnotateEntryFunc()​

如果它是 IRModule 中唯一的函数，则将 PrimFunc 设置为入口点。

tvm.tir.transform.Apply(ftransform)​

将 ftransform 应用于模块中的每个函数。

此函数是 tvm.tir.transform.prim_func_pass 的薄包装器。

  • 参数：ftransform (tvm.tir.PrimFunc -> tvm.tir.PrimFunc)：转换过程。
  • 返回：fpass：结果通过。
  • 返回类型：tvm.transform.Pass。

tvm.tir.transform.BF16ComputeLegalize()​

使 bf16 计算操作合法化。

tvm.tir.transform.BF16StorageLegalize()​

将 bf16 存储类型合法化为 u16。

tvm.tir.transform.BindTarget(target)​

使用给定的目标注释 PrimFunc:param target: target:type target: tvm.target.Target。

tvm.tir.transform.CombineContextCall()​

在宿主函数中组合上下文调用。

tvm.tir.transform.CommonSubexprElimTIR(enable_cse_tir:bool= True, identify_equiv_terms:bool= False)​

用新变量代替冗余计算。

tvm.tir.transform.CompactBufferAllocation(is_strict:bool= True)​

压缩缓冲区访问区域。通过删除未访问的缓冲区，即缩小缓冲区形状并在必要时调整访问区域。

示例

在窄化（narrowing）之前，B 是一个 [16, 16] 的缓冲区，但只访问了一个窄向量 B[i, 0:16]。

for i in range(0, 16):
    with T.block():
        B = T.alloc_buffer(16, 16)
        for j in range(0, 16):
            B[i, j] = A[i, j] + 1
        for j in range(0, 16):
            C[i, j] = B[i, j] + 1

这个 pass 会缩小缓冲区的形状，并相应调整其访问区域。在这个特定例子中，由于只访问了 B 的一个 1 × 16 向量，该 pass 会将 B 缩小为形状 [1, 16]，并将访问 B[i, j] 改为 B[0, j]。

for i in range(0, 16):
    with T.block():
        B = T.alloc_buffer(1, 16)
        for j in range(0, 16):
            B[0, j] = A[i, j] + 1
        for j in range(0, 16):
            C[i, j] = B[0, j] + 1
  • 参数：is_strict (bool)：确保压缩后的形状始终小于原始形状。否则，允许增大形状以匹配实际访问的缓冲区。
  • 返回：fpass：结果通过。
  • 返回类型：tvm.transform.Pass。

tvm.tir.transform.ConvertBlocksToOpaque()​

将所有块变量替换为它们绑定的 PrimExprs，由 BlockRealize 中的相应 iter_values 指示，然后通过删除 BlockRealize 中的所有 iter_values 和 Block 中的 iter_vars 将块转换为不透明的块。

tvm.tir.transform.ConvertForLoopsToSerial()​

将并行 For 循环转换为串行 For 循环。

tvm.tir.transform.ConvertSSA()​

将 IRModule 转换为 SSA 形式。

此过程处理同一 tir.Var 出现在同一模块内多个函数中的情况。例如，将一个函数中的片段提取到另一个函数后，同一个 tir.Var 可能既在原始函数的主函数中定义，又在提升函数中作为参数定义。

tvm.tir.transform.DecorateDeviceScope()​

将所有函数体装饰为设备函数。

tvm.tir.transform.DefaultGPUSchedule()​

该过程为 PrimFuncs 设置默认线程绑定，包括符号形状函数，允许它们在 GPU 设备上构建和执行。它检查 PrimFunc 内的所有块，并根据循环范围和目标信息（例如最大线程块数和每个块的最大线程数）进行循环融合、拆分和重新排序操作。

此过程的主要目标并非优化性能，而是为非调度或符号形状的 PrimFuncs 生成有效的 GPU 内核。此过程目前仅适用于 CUDA 目标平台。

tvm.tir.transform.ExtractPrimFuncConstants()​

收集并统一非标量常量到模块的属性「常量」数组。

tvm.tir.transform.FP8ComputeLegalize(promote_dtype_str:str= 'float32')​

使 fp8 计算操作合法化。

  • 参数：promote_dtype (str)：我们将 fp8 提升到的数据类型，选项：float16/float32。
  • 返回：fpass：结果通过。
  • 返回类型：tvm.transform.Pass。

tvm.tir.transform.FP8StorageLegalize()​

将 fp8 存储类型合法化为 u8。

tvm.tir.transform.Filter(fcond:Callable)​

过滤掉不满足给定条件的 PrimFuncs。fcond 应该是一个接受 primfunc 并返回布尔值的函数。

tvm.tir.transform.FlattenBuffer()​

对于不包含不透明块的 TIR，将多维 BufferLoad 和 BufferStore 展平为单维 BufferLoad/BufferStore。

tvm.tir.transform.ForceNarrowIndexToInt32()​

强制将索引表达式和整数缓冲区缩小为 int32 dtype。

在默认情况下不应使用此通行证。

tvm.tir.transform.HoistExpression()​

HoistIfThenElse 的通用版本。

将循环不变表达式提升到符合条件的循环之外。搜索以下类型的表达式：

  • LetStmt 绑定。
  • IfThenElse 条件。
  • 布尔运算符。
  • 返回：fpass：结果通过。
  • 返回类型：tvm.transform.Pass。

classtvm.tir.transform.HoistExpressionConfig​

提升表达式传递的配置。

propertyhoisted_conditionals​

要提升的布尔表达式类型的位标志。

propertyhoisted_let_bindings​

要提升的 let 绑定类型的位标志。

tvm.tir.transform.HoistIfThenElse(variant:str|None= None)​

将循环不变的 IfThenElse 节点提升到符合条件的循环之外。

  • 参数：variant (Optional*[***String])：pass 的变体。variant 可以具有以下任意一个值 [“basic”, None(Default)]。基本变体支持基本提升场景，其中它期望 For 和 If 节点连续到位，并且不涉及全局范围变量或更高级的场景。默认变体支持所有提升场景，即 {“Basic” + “Advanced”} 通过 PassContext 配置控制支持，如下所示：

config={“tir.HoistIfThenElse”: {“support_block_scope_hoisting”: True}}。

classtvm.tir.transform.HoistIfThenElseConfig​

提升机配置 if then else pass。

propertysupport_block_scope_hoisting​

提升 if cond 与块作用域变量。

class tvm.tir.transform.HoistedConditionals(value)​

HoistExpressionConfig.conditional_types 中使用的标志。

每个位标志代表一种应提升到尽可能最外层循环的表达式类型。

Never = 0​

不提升条件。

IfElseStmt = 1​

如果设置，则在 IfElseStmt 中寻找提升候选。

IfElseExpr = 2​

如果设置，则在 tir.if_then_else 中寻找提升候选。

BooleanExpression = 4​

如果设置，则在所有布尔表达式中寻找提升候选。

UsingBlockVar = 8​

如果设置，则允许提升使用块变量的条件（例如 threadIdx.x）。

All = 15​

启用所有条件提升。

class tvm.tir.transform.HoistedLetBindings(value)​

HoistExpressionConfig.let_binding_types 中使用的标志。

每个位标志代表一种 let 绑定表达式，应将其提升到尽可能最外层的循环。

Never = 0​

不提升 let 绑定。

RequiredByConditional = 1​

提升条件使用的绑定。

LetStmt = 2​

LetStmt 中发生的绑定。

LetExpr = 4​

Let 表达式中发生的绑定。

All = 7​

启用所有 let 绑定的提升。

tvm.tir.transform.InferFragment()​

使用张量内在函数推断 TensorCore 片段信息。

tvm.tir.transform.InjectDoubleBuffer()​

注入双缓冲区语句。

classtvm.tir.transform.InjectDoubleBufferConfig​

注入双缓冲区通道的配置。

propertysplit_loop​

Split loop factors  分割环系数。

tvm.tir.transform.InjectPTXAsyncCopy()​

使用异步复制将全局到 CUDA 上的共享内存复制重写。

tvm.tir.transform.InjectPTXLDG32(enable_inject_ptx_intrin:bool= True)​

注入 ptx.ldg.32 内部函数。

  • 参数：enable_inject_ptx_intrin (bool)：如果为 True，则注入 ptx.ldg.32 内部函数。

tvm.tir.transform.InjectPermutedLayout()​

在 mma 中注入排列布局。

tvm.tir.transform.InjectRollingBuffer()​

注入滚动缓冲区语句。

tvm.tir.transform.InjectSoftwarePipeline()​

将带注释的循环转换为并行生产者和消费者的流水线循环。

tvm.tir.transform.InjectVirtualThread()​

注入虚拟线程循环。

tvm.tir.transform.InlinePrivateFunctions()​

内联调用私有函数。

tvm.tir.transform.InstrumentBoundCheckers()​

instrument 绑定检查器。

tvm.tir.transform.InstrumentProfileIntrinsics()​

插入对 instrument 函数和循环级分析的内部调用。

tvm.tir.transform.LiftThreadBinding()​

将相同的线程绑定提升到其 LCA 环。

tvm.tir.transform.LoopPartition()​

注入虚拟线程循环。

classtvm.tir.transform.LoopPartitionConfig​

配置循环分区传递。

propertyno_unroll_loop_with_extent_one​

不要展开范围为 1 的循环。

propertypartition_const_loop​

Split constant loop  拆分常量循环。

propertyunroll_loop_with_partition_hint_no_interval​

展开循环，有 pragma_loop_partition_hint 且无间隔。

tvm.tir.transform.LowerAsyncDMA()​

降低异步 DMA 至 DMA。

tvm.tir.transform.LowerAutoCopy()​

自动对自动复制块进行内存优化。

tvm.tir.transform.LowerCrossThreadReduction()​

降低从线程绑定到内部函数调用的跨线程减少。

tvm.tir.transform.LowerCustomDatatypes()​

降低自定义数据类型。

有关添加自定义数据类型的更多信息，请参阅 tvm::datatypes::Registry。

tvm.tir.transform.LowerDeviceKernelLaunch()​

降低跨设备函数调用。

在此过程之前，主机到设备的调用表示为子程序调用，并在内部指定环境参数（例如 env_thread）。设备函数是内部函数，没有 tvm::attr::kGlobalSymbol 属性。

经过此过程后，主机到设备的调用将表示为内置的 tvm_call_packed 函数。设备函数是一个对外暴露的函数，具有非空的 tvm::attr::kGlobalSymbol 属性。

tvm.tir.transform.LowerDeviceStorageAccessInfo()​

降低设备上附加存储的访问信息。

所有存储访问分析完成后运行此过程。

tvm.tir.transform.LowerInitBlock()​

将块初始化语句转换为 IfThenElse 语句 。

tvm.tir.transform.LowerIntrin()​

降低目标特定的内在调用。

tvm.tir.transform.LowerMatchBuffer()​

删除块内的匹配缓冲区。此外，它将验证绑定。

tvm.tir.transform.LowerOpaqueBlock()​

解除封锁以确保 TIR 不会再次被安排。

tvm.tir.transform.LowerTVMBuiltin()​

降低 TVM 内置内在函数。

tvm.tir.transform.LowerThreadAllreduce()​

降低跨线程 alleduce。

tvm.tir.transform.LowerVtcmAlloc()​

降低 vtcm 分配。

tvm.tir.transform.LowerWarpMemory()​

降低 warp 内存访问到低级设备相关的函数调用。

tvm.tir.transform.MakePackedAPI()​

将模块中的 PrimFuncs 转换为打包的 func API。

在此过程之前，PrimFunc 可能已在 PrimFuncNode::buffer_map 中定义了 Buffer 参数。此过程将使用 buffer_map，并使用它生成实现基于打包的 TVM FFI API 的参数。

对于静态形状，BufferNode::shape、BufferNode::strides 和 BufferNode::elem_offset 成员变量用于对用户提供的 DLTensor*或 tvm.nd.array 参数中相应成员变量生成运行时检查。（例如，接受形状为[16,32]的缓冲区的 PrimFunc 验证 DLTensor::shape 数组是否为[16,32]。）。

对于动态缓冲区，其中一个或多个 BufferNode 成员变量使用未由其他 PrimFunc 参数定义的 tir.Var ，这些变量将用于根据相应的 DLTensor 成员定义变量。（例如，接受形状为[tir.Var(“n”), tir.Var(“m”)]的缓冲区的 PrimFunc ，在传递 形状为[16,32]的 DLTensor 时，将根据参数的形状定义n = 16和 n=32 。

tvm.tir.transform.MakeUnpackedAPI()​

将模块中的 PrimFuncs 转换为与内部调用兼容的C API。

在此过程之前，PrimFunc 可能已在 PrimFuncNode::buffer_map 中定义了 Buffer 参数。此过程将使用 buffer_map，并使用它生成可由 C API 直接调用的 T参数（例如 float32 ）。

对于静态形状，无需执行运行时验证来确认参数缓冲区的形状是否与预期形状匹配。对于动态形状，MakeUnpackedAPI 要求将动态参数作为单独的 tir.Var 参数传递。

tvm.tir.transform.ManifestSharedMemoryLocalStage()​

为 GPU 上的共享内存访问添加显式本地阶段。

tvm.tir.transform.MergeSharedMemoryAllocations()​

此过程将多个 TIR 级共享内存分配合并为一个分配。

tvm.tir.transform.NarrowDataType(target_bits:int)​

将 stmt 中的 PrimExpr 数据类型缩小到 target_bits。

  • 参数：
    • target_bits (int)：目标位配置。
  • 返回：fpass：结果通过。
  • 返回类型：tvm.transform.Pass。

在 FlattenBuffer 之后运行此过程。

tvm.tir.transform.PlanAndUpdateBufferAllocationLocation()​

将缓冲区分配定位到精确位置（通常是缓冲区访问的 lca）。此过程将在分配位置注入带有 alloc_buffers 的不透明块。

tvm.tir.transform.PointerValueTypeRewrite()​

重写参数的指针内容类型，以及函数内部的 Alloc，以使用最常访问的类型进行加载/存储，从而尽可能避免在后端进行指针转换。

tvm.tir.transform.ReduceBranchingThroughOvercompute()​

通过引入过度计算来减少分支。

classtvm.tir.transform.ReduceBranchingThroughOvercomputeConfig​

通过超额计算传递减少分支的配置。

propertyuse_dataflow_analysis​

如果为 true，则传播已知缓冲区值并用于静态证明超额计算是有效的。

tvm.tir.transform.RemoveAssume()​

删除所有 builtin::assume 实例。

tvm.tir.transform.RemoveNoOp()​

从 Stmt 中删除 No Op。

classtvm.tir.transform.RemoveNoOpConfig​

删除无作通道的配置。

propertymax_simplification_steps​

如果不为零，则 RewriteSimplifier 将在指定的步骤数后引发错误。用于调试和测试目的。

 propertyuse_dataflow_analysis​

如果为 true，则传播已知的缓冲区值并用于静态证明语句为 no-ops。

tvm.tir.transform.RemoveStoreUndef()​

从 Stmt 中删除未定义值的存储。

tvm.tir.transform.RemoveWeightLayoutRewriteBlock(skip_ndarray_rewrite=False)​

在调整阶段进行基准测试之前删除权重布局重写块。

  • 参数：
    • skip_ndarray_rewrite (bool) ：

如果为 True，则将跳过根据给定索引图精确重写 NDArray 的操作。仅正确转换 NDArray 的形状，并用随机值填充目标数组的内容。在 MetaSchedule 调优过程中，如果多次调用此过程，则重写前后 NDArray 的原始数据无关紧要。由于使用 IndexMap 的 MapNDArray 进行 NDArray 布局重写目前速度较慢，因此有时需要跳过精确的重写。

tvm.tir.transform.RenormalizeSplitPattern()​

将分割模式从 floordiv(floormod()) 重新规范化为 floormod(floordiv())。

tvm.tir.transform.RewriteUnsafeSelect()​

检测并重写包含内存访问的不安全选择。

tvm.tir.transform.Simplify()​

对语句和表达式进行算术简化。

classtvm.tir.transform.SimplifyConfig​

简化传递的配置。

propertyapply_constraints_to_boolean_branches​

如果为 true，则在另一个分支提供的约束下简化 AND/OR 的每个分支。

 propertyconvert_boolean_to_and_of_ors​

如果为 true，则将条件简化为 OR 的 AND。

 propertypropagate_knowns_to_prove_conditional​

如果为 true，则传播已知缓冲区值并用于静态证明条件。

 propertypropagate_knowns_to_simplify_expressions​

如果为 true，则传播已知缓冲区值，并尽可能用于替换 BufferLoad。

 propertytransitively_prove_inequalities​

如果为 true，则使用作用域约束的传递组合简化条件。

tvm.tir.transform.SkipAssert()​

跳过断言语句。

tvm.tir.transform.SplitHostDevice()​

将函数分为主机函数和设备函数。

tvm.tir.transform.StorageRewrite()​

重写存储分配模式。

将分配移至尽可能最外层的作用域。尝试在分配之间共享空间，以便在可能的情况下制定静态分配计划。

tvm.tir.transform.ThreadSync(storage_scope:str)​

在共享缓冲区的并行读/写之间插入同步。

  • 参数：storage_scope (str)：目标存储范围。
  • 返回：fpass：结果通过。
  • 返回类型：tvm.transform.Pass。

tvm.tir.transform.TransformMmaBufferLayout()​

变换 mma 缓冲区布局。

tvm.tir.transform.UnifyThreadBinding()​

统一所有“blockIdx.x/y/z”、“threadIdx.x/y/z”和“vthread.x/y/z”的线程绑定。统一之前，绑定到同一线程轴的两个变量（例如“threadIdx.x”）在其 AttrStmts 中使用不同的 IterVar 和变量。统一之后，我们将为它们使用合并后的 IterVar 和变量。

vthread 是一种将被弃用的遗留行为，但 vthread 的线程绑定在本阶段仍然统一。请使用 vthread.x、vthread.y 和 vthread.z 代替。

tvm.tir.transform.UnrollLoop()​

展开以 unroll 标记的常量循环。

此过程还自动将 pragma unroll 标签附加到符合标准的循环。

classtvm.tir.transform.UnrollLoopConfig​

展开循环通道的配置。

propertyauto_max_depth​

可以自动展开的循环的最大嵌套级别。

 propertyauto_max_extent​

将展开的循环的最大范围。

 propertyauto_max_step​

循环中要自动展开的步数阈值。

 propertyexplicit_unroll​

是否显式展开循环而不是设置编译指示。

 propertyunroll_local_access​

是否始终展开本地访问。

tvm.tir.transform.UseAssumeToReduceBranches()​

此过程尝试通过过度计算填充区域的值来消除布局特定的填充分支。消除分支将有助于矢量化代码，并提升元素级操作的性能。

tvm.tir.transform.VectorizeLoop(enable_vectorize:bool= True)​

降低矢量化循环。

  • 参数：enable_vectorize (bool)：是否启用矢量化。关闭时将降低到标量循环。
  • 返回：fpass：结果通过。
  • 返回类型：tvm.transform.Pass。

tvm.tir.transform.VerifyMemory()​

验证 func 是否包含非法主机端直接内存访问。

tvm.tir.transform.VerifyVTCMLimit(limit=None)​

验证分配的 vtcm 内存的大小是否满足限制。 *

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