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。
当提供 pass_func 时，此函数返回一个回调。否则，它使用给定的优化函数返回创建的函数 pass。
- 参数:
pass_func (Optional[Callable[(tvm.tir.PrimFunc, IRModule, PassContext) -> tvm.tir.PrimFunc]]) – 转换函数或类。
opt_level (int) – 此模块 pass 的优化级别。
name (Optional[str]) – 函数 pass 的名称。名称可以为空。在这种情况下，优化函数的名称将用作 pass 名称。
required (Optional[List[str]]) – 函数 pass 依赖的 pass 列表。
- 返回:
create_function_pass – 如果未提供 pass_func，则返回装饰器，否则返回装饰后的结果。返回的装饰器根据输入具有两种行为：当我们装饰一个 pass 函数时，将返回一个新的 FunctionPass。当我们装饰一个类类型时，将返回一个新的 FunctionPass 类。
- 返回类型:
Union[Callable, FunctionPass]
示例
以下代码块装饰了一个函数 pass 类。
@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): # just for demo purposes # transform func to new_func return self.new_func
以下代码通过装饰用户定义的转换函数来创建一个函数 pass。
@tvm.tir.transform.prim_func_pass(opt_level=2) def transform(func, mod, ctx): # my transformations here. return func function_pass = transform assert isinstance(function_pass, transform.FunctionPass) assert function_pass.info.opt_level == 2 # Given a module m, the optimization could be invoked as the following: updated_mod = function_pass(m) # Now constant folding should have been applied to every function in # the provided module m. And the updated module will be returned.
- class tvm.tir.transform.PrimFuncPass
一个在模块中的每个
tvm.tir.PrimFunc()上工作的 pass。函数 pass 类应该通过 py:func:tvm.tir.transform.function_pass 创建。
- tvm.tir.transform.AnnotateDeviceRegions()
注解应该在设备上运行的位置
插入 AttrStmt 节点，指定 PrimFunc 内的区域应该在其上执行的目标。仅修改具有 tvm::attr::kTarget 属性并且该目标定义了主机的函数。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.AnnotateEntryFunc()
如果 PrimFunc 是 IRModule 中唯一的函数，则将其设置为入口点。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.Apply(ftransform)
将 ftransform 应用于模块中的每个函数。
此函数是 tvm.tir.transform.prim_func_pass 的一个薄封装
- 参数:
ftransform (tvm.tir.PrimFunc -> tvm.tir.PrimFunc) – 转换 pass。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.BF16ComputeLegalize()
合法化 bf16 计算 Ops。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.BF16StorageLegalize()
将 bf16 存储类型合法化为 u16。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.BindTarget(target)
使用给定的目标注解 PrimFunc。 :param target: target :type target: tvm.target.Target
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.CombineContextCall()
合并主机函数中的上下文调用。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.CommonSubexprElimTIR(enable_cse_tir: bool = True, identify_equiv_terms: bool = False)
用新变量替换冗余计算。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.CompactBufferAllocation(is_strict: bool = True)
压缩缓冲区访问区域。通过删除未访问的缓冲区区域来压缩，即缩小缓冲区形状并在必要时调整访问区域。
示例
在缩小之前，
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 – 结果 pass
- 返回类型:
- tvm.tir.transform.ConvertBlocksToOpaque()
将所有块变量替换为它们绑定的 PrimExpr，由 BlockRealize 中的相应 iter_values 指示，然后通过删除 BlockRealize 中的所有 iter_values 和 Block 中的 iter_vars 将块转换为不透明块。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.ConvertForLoopsToSerial()
将并行 For 循环转换为串行 For 循环。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.ConvertSSA()
将 IRModule 转换为 SSA 形式。
此 pass 处理在同一模块内的多个函数中出现相同 tir.Var 的情况。例如，在将片段从一个函数提取到另一个函数之后，相同的 tir.Var 可能在原始函数的正文中定义，并在提升的函数中作为参数定义。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.DecorateDeviceScope()
将所有函数的主体装饰为设备函数。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.DefaultGPUSchedule()
此 pass 为 PrimFunc 设置默认线程绑定，包括符号形状函数，允许它们在 GPU 设备上构建和执行。它检查 PrimFunc 中的所有块，并根据循环范围和目标信息（例如最大线程块数和每个块的最大线程数）执行循环融合、拆分和重新排序操作。
此 pass 的主要目标不是优化性能，而是为未调度或符号形状 PrimFunc 生成有效的 GPU 内核。该 pass 目前仅适用于 CUDA 目标。
- 返回:
ret
- 返回类型:
- tvm.tir.transform.ExtractPrimFuncConstants()
收集并将 tir 非标量常量统一到模块的属性 ‘Constants’ 数组。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.FP8ComputeLegalize(promote_dtype_str: str = 'float32')
合法化 fp8 计算 Ops。
- 参数:
promote_dtype (str) – 我们将 fp8 提升到的数据类型，选项：float16/float32。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.FP8StorageLegalize()
将 fp8 存储类型合法化为 u8。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.Filter(fcond: Callable)
过滤掉不满足给定条件的 PrimFunc。fcond 应该是一个接受 primfunc 并返回布尔值的函数。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.FlattenBuffer()
对于不包含不透明块的 TIR，将多维 BufferLoad 和 BufferStore 展平为单维 BufferLoad/BufferStore。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.ForceNarrowIndexToInt32()
强制将索引表达式和整数缓冲区缩小为 int32 数据类型。
- 返回:
fpass – 结果 pass
- 返回类型:
注意
此 pass 不应在默认情况下使用。
- tvm.tir.transform.HoistExpression()
HoistIfThenElse 的通用版本。
将循环不变表达式提升到符合条件的循环之外。搜索以下表达式：
LetStmt 绑定
IfThenElse 条件
布尔运算符
- 返回:
fpass – 结果 pass
- 返回类型:
- 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_hosting”: True}}
- 返回:
fpass – 结果 pass
- 返回类型:
- 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 片段信息。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.InjectDoubleBuffer()
注入双缓冲语句。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.InjectPTXAsyncCopy()
在 CUDA 上使用异步复制重写全局到共享内存的复制。
- 返回:
fpass – 结果 pass
- 返回类型:
- 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 中注入置换布局
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.InjectRollingBuffer()
注入滚动缓冲区语句。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.InjectSoftwarePipeline()
将注解的循环转换为流水线循环，以并行化生产者和消费者
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.InjectVirtualThread()
注入虚拟线程循环。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.InlinePrivateFunctions()
内联对私有函数的调用
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.InstrumentBoundCheckers()
检测边界检查器。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.InstrumentProfileIntrinsics()
插入内在函数调用以检测函数和循环级别的性能分析。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LegalizePackedCalls()
合法化 packed 调用，使其参数包装在 TVMValue 中
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LiftThreadBinding()
将相同的线程绑定提升到其 LCA 循环。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LoopPartition()
注入虚拟线程循环。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerAsyncDMA()
将异步 DMA 降低为 DMA。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerAutoCopy()
为自动复制块自动执行内存优化
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerCrossThreadReduction()
将跨线程归约从线程绑定降低到内在函数调用。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerCustomDatatypes()
降低自定义数据类型。
有关添加自定义数据类型的更多信息，请参见 tvm::datatypes::Registry。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerDeviceKernelLaunch()
降低跨设备函数调用。
在此 pass 之前，主机到设备调用表示为子例程调用，并在内部指定环境参数（例如 env_thread）。设备函数是一个内部函数，没有 tvm::attr::kGlobalSymbol 属性。
在此 pass 之后，主机到设备调用表示为 tvm_call_packed 内置函数。设备函数是一个外部公开的函数，具有非空的 tvm::attr::kGlobalSymbol 属性。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerDeviceStorageAccessInfo()
降低设备上附加的存储访问信息。
- 返回:
fpass – 结果 pass
- 返回类型:
注意
在所有存储访问分析完成后运行此 pass。
- tvm.tir.transform.LowerInitBlock()
将块初始化语句降级为 IfThenElse 语句。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerIntrin()
降低目标特定的内联函数调用。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerMatchBuffer()
移除块内的匹配缓冲区。此外，它还会验证绑定。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerOpaqueBlock()
移除块以确保 TIR 不能再次被调度。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerTVMBuiltin()
降低 tvm 内置内联函数。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerThreadAllreduce()
降低跨线程规约 (allreduce)。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerVtcmAlloc()
降低 vtcm 分配。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.LowerWarpMemory()
降低 warp 内存访问为底层设备相关函数调用。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.MakePackedAPI()
将模块中的 PrimFuncs 转换为 packed func API。
在此 pass 之前，PrimFunc 可能在 PrimFuncNode::buffer_map 中定义了 Buffer 参数。此 pass 消耗 buffer_map，并使用它来生成实现 PackedFunc API 的 TVMArgs 和 TVMRetValue* 参数。
对于静态形状，BufferNode::shape、BufferNode::strides 和 BufferNode::elem_offset 成员变量用于在用户提供的 DLTensor* 或 tvm.nd.array 参数中生成对相应成员变量的运行时检查。（例如，接受形状为 [16,32] 的缓冲区的 PrimFunc 会验证 DLTensor::shape 数组是否为 [16,32]。）
对于动态 Buffer，其中一个或多个 BufferNode 成员变量使用未被其他 PrimFunc 参数定义的 tir.Var，这些变量将用于根据相应的 DLTensor 成员定义变量。（例如，接受形状为 [tir.Var(“n”), tir.Var(“m”)] 的缓冲区的 PrimFunc，当传递形状为 [16,32] 的 DLTensor 时，将根据参数的形状定义 n = 16 和 n=32。）
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.MakeUnpackedAPI()
将模块中的 PrimFuncs 转换为与内部调用兼容的 C API。
在此 pass 之前，PrimFunc 可能在 PrimFuncNode::buffer_map 中定义了 Buffer 参数。此 pass 消耗 buffer_map，并使用它来生成可以被 C API 直接调用的 T* 参数（例如 float32*）。
对于静态形状，不执行运行时验证来确认参数缓冲区的形状是否与预期形状匹配。对于动态形状，MakeUnpackedAPI 要求动态参数作为单独的 tir.Var 参数传递。
- 返回:
fpass – 结果 pass
- 返回类型:
为 GPU 上的共享内存访问添加显式的本地阶段。
- 返回:
fpass – 结果 pass
- 返回类型:
此 pass 将多个 TIR 级别的共享内存分配合并为一个分配。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.NarrowDataType(target_bits: int)
将 stmt 中的 PrimExpr 数据类型缩小到 target_bits。
- 参数:
target_bits (int) – 目标位配置。
- 返回:
fpass – 结果 pass
- 返回类型:
注意
在 FlattenBuffer 之后运行此 pass。
- tvm.tir.transform.PlanAndUpdateBufferAllocationLocation()
将缓冲区分配定位到精确位置（通常是缓冲区访问的 lca）。此 pass 将在分配站点注入带有 alloc_buffers 的 opaque block。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.PointerValueTypeRewrite()
重写参数的指针内容类型，以及函数内部的 Alloc，以使用最频繁访问的类型进行加载/存储，以尽可能避免后端中的指针转换。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.ReduceBranchingThroughOvercompute()
通过引入过度计算来减少分支。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.RemoveAssume()
移除所有 builtin::assume 的实例。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.RemoveNoOp()
从 Stmt 中移除 No Op。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.RemoveStoreUndef()
从 Stmt 中移除未定义值的存储。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.RemoveWeightLayoutRewriteBlock(skip_ndarray_rewrite=False)
在 tuning 阶段的基准测试之前，移除权重布局重写块。
- 参数:
skip_ndarray_rewrite (bool) –
如果为 True，则会跳过根据给定索引映射精确重写 NDArray。仅正确转换 NDArray 的形状，并且目标数组的内容将填充随机值。
当在 MetaSchedule tuning 期间多次调用此 pass 时，NDArray 的原始数据在重写前后无关紧要。由于使用 IndexMap 的 MapNDArray 进行 NDArray 布局重写目前速度较慢，因此有时需要跳过精确重写。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.RenormalizeSplitPattern()
将拆分模式从 floordiv(floormod()) 重新归一化为 floormod(floordiv())。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.RewriteUnsafeSelect()
检测并重写包含内存访问的不安全 select。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.Simplify()
对语句和表达式运行算术简化。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.SkipAssert()
跳过 assert 语句。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.SplitHostDevice()
将函数拆分为主机函数和设备函数。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.StorageRewrite()
重写存储分配模式。
将分配移动到最外层的可能作用域。尝试在分配之间共享空间，以便在可能的情况下制定静态分配计划。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.ThreadSync(storage_scope: str)
在共享缓冲区的并行读/写之间插入同步。
- 参数:
storage_scope (str) – 目标存储作用域。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.TransformMmaBufferLayout()
转换 mma 缓冲区布局。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.UnifyThreadBinding()
统一 “blockIdx.x/y/z”、“threadIdx.x/y/z” 和 “vthread.x/y/z” 的所有线程绑定。在统一之前，绑定到线程轴（例如 “threadIdx.x”）的两个变量在其 AttrStmt 中使用不同的 IterVar 和变量。统一后，我们为它们使用一个统一的 IterVar 和一个变量。
- 返回:
fpass – 结果 pass
- 返回类型:
注意
vthread 是一种将被弃用的遗留行为，尽管 vthread 的线程绑定仍然在此 pass 中统一。请改用 vthread.x、vthread.y 和 vthread.z。
- tvm.tir.transform.UnrollLoop()
展开由 unroll 标记的常量循环。
此 pass 还会自动将 pragma unroll 标签附加到符合标准的循环。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.UseAssumeToReduceBranches()
此 pass 尝试通过过度计算填充区域的值来消除特定于布局的填充分支。消除分支将有助于向量化代码，并提高元素级操作的性能。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.VectorizeLoop(enable_vectorize: bool = True)
降低向量化循环。
- 参数:
enable_vectorize (bool) – 是否启用向量化。关闭时将降低为标量循环。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.VerifyMemory()
验证 func 是否包含非法的主机端直接内存访问。
- 返回:
fpass – 结果 pass
- 返回类型:
- tvm.tir.transform.VerifyVTCMLimit(limit=None)
验证分配的 vtcm 内存大小是否满足限制。
- 返回:
fpass – 结果 pass
- 返回类型: