尧图网站建设 
TVM s_tir.dlight 模块完全指南开箱即用的默认调度规则体系与 GPU/CPU/Adreno 实战【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址: https://gitcode.com/gh_mirrors/tv/tvm导读tvm.s_tir.dlight是 TVM 中面向深度学习负载的开箱即用调度库它内置了一批针对 GEMM、GEMV、Reduction、Transpose、RMSNorm 等典型算子的调度规则ScheduleRule无需任何调优即可为 IRModule 中的每个PrimFunc生成高质量的默认调度也可以为 MetaSchedule 提供可调优的搜索空间。本文以 docs/reference/api/python/s_tir/dlight.rst 为骨架结合仓库中 python/tvm/s_tir/dlight 的源码实现与 Relax 后端集成路径系统讲解该模块的公共 API、各子模块职责、调度规则的运行原理以及如何通过ApplyDefaultSchedule接入自定义编译管线。一、模块定位DLight 是什么解决什么问题DLightDeep Learning Lightweight Scheduler的核心目标正如其包级 docstring 所言provide efficient schedules out-of-box for deep learning workloadspython/tvm/s_tir/dlight/init.py。它在传统两条路线之间取了一个平衡点完全手工写 schedule需要开发者对每个算子和每种硬件写大量T.schedule指令成本高、难维护MetaSchedule 自动调优质量上限高但需要 builder/runner 实测、代价大不适合推理部署这类要求快速出结果的场景。DLight 的做法是把常见算子模式的调度经验固化成规则rule在编译期对 TIR 层面做结构分析直接命中规则并产出完整 schedule从而做到零调优、开箱即用。从文档的 RST 结构看tvm.s_tir.dlight的公开面被组织成七个部分docs/reference/api/python/s_tir/dlight.rstRST 章节对应包路径职责tvm.s_tir.dlightpython/tvm/s_tir/dlight/init.py顶层统一入口聚合导出规则类、分析工具与基础设施tvm.s_tir.dlight.gpupython/tvm/s_tir/dlight/gpu/init.pyGPU 通用调度规则CUDA/ROCm/Vulkan/Metal 均适用tvm.s_tir.dlight.adrenopython/tvm/s_tir/dlight/adreno/init.pyAdreno 移动 GPU 专用规则卷积/池化/布局转换tvm.s_tir.dlight.cpupython/tvm/s_tir/dlight/cpu/init.pyCPU 通用规则GEMV、Reductiontvm.s_tir.dlight.analysispython/tvm/s_tir/dlight/analysis调度前结构分析工具SBlockInfo、IterInfo 等tvm.s_tir.dlight.basepython/tvm/s_tir/dlight/base基础设施ScheduleRule 抽象、ApplyDefaultSchedule pass、公共工具函数tvm.s_tir.dlight.benchmarkpython/tvm/s_tir/dlight/benchmark调度结果提取与基准测试辅助注gpu子模块在 docstring 中特别说明For CUDA/ROCm/Vulkan/Metal-specific rules, usetvm.s_tir.dlight.cuda/rocm/vulkan/metalinstead即这里的规则是GPU 通用而非某厂商私有的。二、顶层公共 APItvm.s_tir.dlight导出内容顶层__init__.pypython/tvm/s_tir/dlight/init.py共导出三类对象后端子模块gpu、adreno、cpu分别承载对应硬件的规则集合分析工具来自dlight.analysisSBlockInfo对 PrimFunc 的调度块SBlock做结构信息封装IterInfo描述单个循环迭代变量的信息是否归约、是否空间等normalize_prim_func把函数归一化为便于规则匹配的标准形态基础设施来自dlight.baseApplyDefaultSchedule一个module_pass的 IRModule pass把一组规则应用到模块内所有 PrimFunc 上ScheduleRule规则的抽象基类try_inline/try_inline_contiguous_spatial通用的 producer/consumer 内联辅助函数。因此用户最常见的导入方式就是from tvm.s_tir import dlight as dl # 使用 GPU 通用规则 dl.gpu.Matmul() # 使用 CPU 规则 dl.cpu.Reduction() # 使用分析工具 dl.SBlockInfo(func) # 使用默认调度 pass dl.ApplyDefaultSchedule(dl.gpu.Matmul(), dl.gpu.Fallback())三、ScheduleRule抽象规则的定义与生命周期ScheduleRule定义在 python/tvm/s_tir/dlight/base/schedule_rule.py是整个 DLight 体系最核心的接口。它的语义是A thin wrapper on an arbitrary function that can be used to schedule a TIR PrimFunc.3.1apply方法核心契约每个规则必须实现apply(func, target, tunable)functirx.PrimFunc待调度的 TIR 函数targettvm.target.Target规则为其生成 schedule 的编译目标tunable布尔标志控制规则产出的形式tunableFalse只能返回一个s_tir.Schedule且其中不允许包含可调优指令tunable instructions——这是默认调度模式的约束tunableTrue可以返回一个或一组s_tir.Schedule且允许包含可调优指令——这是为 MetaSchedule 提供搜索空间的模式返回值可以是None、单个s_tir.Schedule或list[s_tir.Schedule]。返回None表示当前规则不适用于该 PrimFunc调度引擎会继续尝试下一条规则。3.2from_callable用普通函数快速声明规则ScheduleRule.from_callable(name)是一个装饰器工厂让你不必继承基类就能把一个普通函数包装成规则from tvm.s_tir import dlight as dl from tvm import tirx from tvm.target import Target dl.ScheduleRule.from_callable(MyRule) def my_rule(func: tirx.PrimFunc, target: Target, tunable: bool): # 对 func 做结构判断返回 None 或 s_tir.Schedule ...其内部实现是为传入函数动态生成一个_Rule子类并把函数绑定到apply上同时将类名设置为传入的name源码见 schedule_rule.py。3.3is_target_available默认返回True子类可重写用于在特定 target如无 CUDA 的机器上提前声明规则不可用从而被调度引擎跳过。3.4 规则匹配的调度语义以_apply_rules为例ApplyDefaultSchedule内部通过_apply_rulestransform.py串行遍历规则列表for rule in rules: space rule.apply(func, target, tunable) if space is None: continue # 当前规则不适用尝试下一条 if isinstance(space, s_tir.Schedule): space [space] return space # 命中第一条规则即返回 return None这种顺序尝试、先命中者胜出的机制决定了规则列表的排列顺序很有讲究越通用的兜底规则如Fallback应排在最后越专门的规则如Matmul排在前面。这一点可以从 Relax 的 CUDA 后端默认管线的排列得到印证见下文第六节。四、ApplyDefaultSchedule把规则变成一次 IRModule passApplyDefaultSchedule定义在 python/tvm/s_tir/dlight/base/transform.py是用户接触最多的入口。它是一个module_pass(opt_level0, nameApplyDefaultSchedule)即属于module pass输入输出都是IRModuleopt_level0意味着它在 pass 管线中属于功能性必做而非优化可选类别构造时传入若干ScheduleRule实例ApplyDefaultSchedule(*rules)。4.1 工作流程transform_module对模块内每个函数执行如下逻辑transform.py只处理tirx.PrimFunc且跳过已调度过的函数——判断依据是函数属性tirx.is_scheduled是否为 1_is_scheduled获取目标设备优先读取函数自带的target属性否则回退到Target.current()_get_targetallow_noneFalse表示当前无 target 上下文会直接报错以tunableFalse调用_apply_rules若命中规则将产出的 schedule 对应的新函数写回模块并打上tirx.is_scheduled True标记避免后续被重复调度。4.2 幂等性与去重tirx.is_scheduled标记是这套机制的关键设计它保证了 pass 可重复执行、可与其它 pass 任意组合而不会产生二次调度。自定义编译器如果自己写了调度逻辑同样可以依赖该属性避免与 DLight 冲突。4.3 直接用法示例import tvm from tvm import tirx from tvm.s_tir import dlight as dl from tvm.script import tirx as T T.prim_func def my_gemm(A: T.Buffer((256, 256), float32), B: T.Buffer((256, 256), float32), C: T.Buffer((256, 256), float32)): for i, j, k in T.grid(256, 256, 256): with T.block(C): vi, vj, vk T.axis.remap(SSR, [i, j, k]) with T.init(): C[vi, vj] 0.0 C[vi, vj] A[vi, vk] * B[vk, vj] mod tvm.IRModule({main: my_gemm}) target tvm.target.Target(cuda) # 应用 DLight 默认调度 with target: mod dl.ApplyDefaultSchedule(dl.gpu.Matmul(), dl.gpu.Fallback())(mod) # 之后即可正常走 codegen 生成 CUDA 内核说明上述tunableFalse的默认路径对每个函数只生成一个确定 schedule若想拿到一组候选空间用于自动调优可参考ScheduleRule.apply(tunableTrue)的语义结合 MetaSchedule 的搜索策略使用。五、GPU 规则集tvm.s_tir.dlight.gpugpu子模块python/tvm/s_tir/dlight/gpu/init.py是覆盖度最广的规则集合共导出 8 条规则规则类源文件适用算子模式Matmulgpu/matmul.pyGEMM含 Batch GEMM如全连接、attention 的 QK^T / PVGEMVgpu/gemv.py矩阵-向量乘M 很小如解码阶段的 LinearLowBatchGEMVgpu/low_batch_gemv.pybatch 维度很小的 GEMV进一步特化Reductiongpu/reduction.py规约类算子sum/max/min/softmax 的一部分等GeneralReductiongpu/general_reduction.py一般化规约模式的兜底Transposegpu/transpose.py布局转置算子RMSNormgpu/rmsnorm.pyRMSNorm / LayerNorm 类归一化Fallbackgpu/fallback.py兜底规则上述规则都不命中时保证仍有可用 schedule5.1 源码级理解以Matmul为例Matmul规则gpu/matmul.py是理解整条规则链的绝佳样例其调度流程体现了 DLight 的通用套路结构分析调用get_root_block、IterInfo/SBlockInfo等分析工具识别 GEMM 的三个迭代维度M、N、K区分空间轴与归约轴内联辅助算子利用auto_inline_producers/auto_inline_consumers反复对 producer/consumer 块执行compute_inline直到没有可内联的块为止见 matmul.py把 GEMM 周边 elementwise 算子全部折叠进主计算分块与线程绑定对 M、N 维度做 tile并基于 target 信息做 block/thread 绑定与向量化返回 Schedule生成完整s_tir.Schedule供后续 codegen。auto_inline_*的实现非常直接递归收集所有 producer/consumer 块逐个尝试sch.compute_inline(block)抛异常不可内联就跳过继续下一个直到一轮下来没有任何内联成功才停止。这种能内联就内联的激进策略是 DLight 保持内核简洁高效的基石。5.2Fallback规则的意义Fallback保证任何函数都能被调度因此它是所有 GPU 规则链的最后一道保险。在自定义编译管线中如果你不确定自己的算子集合能被 DLight 覆盖务必把Fallback放在规则列表末尾避免_apply_rules返回None导致整个 pass 对该函数无所作为。六、真实接入案例Relax CUDA 后端如何内置 DLightDLight 并不是一个孤立的工具库它已经被整合进 TVM 的 Relax 编译流程。以 CUDA 后端为例python/tvm/relax/backend/cuda/pipeline.py 中的legalize_passes展示了标准接入姿势def legalize_passes(target: tvm.target.Target): from tvm.s_tir import dlight as dl return [ tvm.relax.transform.LegalizeOps(), tvm.relax.transform.AnnotateTIROpPattern(), tvm.relax.transform.FoldConstant(), tvm.relax.transform.FuseOps(), tvm.relax.transform.FuseTIR(), # 先完成算子融合产出 TIR PrimFunc dl.ApplyDefaultSchedule( # 再对融合后的函数应用 DLight 默认调度 dl.gpu.Matmul(), dl.gpu.GEMV(), dl.gpu.Reduction(), dl.gpu.GeneralReduction(), dl.gpu.Fallback(), ), ]这串代码揭示了 DLight 在完整编译流程中的精确位置FuseOpsFuseTIR先把 Relax 层的算子图融合成一个个 TIRPrimFuncApplyDefaultSchedule紧接着对这些融合函数按Matmul → GEMV → Reduction → GeneralReduction → Fallback的顺序尝试匹配规则之后才进入 dataflow lowering 与最终 codegen见同一文件中的dataflow_lower_passes与finalize_passes。也就是说DLight 取代了传统手工写调度如sch tirx.create_schedule后逐条 bind/vectorize的环节用户只需声明规则列表编译管线自动完成调度。类似的集成还出现在 python/tvm/relax/backend/rocm/pipeline.py、python/tvm/relax/backend/gpu_generic/pipeline.py、python/tvm/relax/backend/adreno/pipeline.py 以及 python/tvm/relax/pipeline.py 中。七、CPU 与 Adreno 规则集7.1 CPUtvm.s_tir.dlight.cpuCPU 规则集python/tvm/s_tir/dlight/cpu/init.py相对精简导出两条CPU 通用规则GEMVcpu/gemv.py矩阵-向量乘重点处理缓存友好的分块与向量化Reductioncpu/reduction.py规约类算子关注寄存器/多线程规约的组织。CPU 场景下 DLight 的核心收益是快速生成多线程 向量化的内核无需手工做parallel/vectorize决策。7.2 Adrenotvm.s_tir.dlight.adrenoAdreno高通移动 GPU规则集python/tvm/s_tir/dlight/adreno/init.py面向移动端推理导出四条规则Conv2dadreno/convolution.py二维卷积移动端最核心的算子Pool2Dadreno/pool.py二维池化LayoutTransformadreno/layout_transform.py布局转换如 NHWC 相关重排Fallbackadreno/fallback.py兜底规则。Adreno 的调度重点关注移动 GPU 的纹理/局部内存特性结合 python/tvm/relax/backend/adreno/pipeline.py 可以看到它在 Relax Adreno 后端管线中同样承担默认调度职责。八、分析工具与基础设施analysis、base、benchmark8.1analysis调度前的结构分析python/tvm/s_tir/dlight/analysis 提供规则共享的分析设施SBlockInfo、IterInfo把 TIR 块与循环迭代结构抽象成规则易于查询的数据结构顶层__init__.py直接再导出见 dlight/init.pycommon_analysis.py常见分析的共享实现gemv.py面向 GEMV 模式的结构识别。规则编写者应尽量复用这些分析工具而不是各自解析 TIR 语法树从而保证整个规则库对结构判断的口径一致。8.2base基础设施与工具函数base包python/tvm/s_tir/dlight/base除了前文讲过的ScheduleRule、ApplyDefaultSchedule、try_inline/try_inline_contiguous_spatial定义于 base/common_schedules.py还导出一组调度参数启发式工具base/utils.pymax_threads_per_block查询目标允许的最大线程数/块大小供规则决定 block 配置suggest_threads_per_block根据循环长度启发式地给出线程数建议auto_vectorize自动选择向量化宽度get_extent/get_bytes读取循环 extent 与数据类型字节数用于计算分块大小。这些函数是 DLight 规则按 target 自适应的秘密武器规则不写死线程数而是借助这些查询函数在编译期动态推导这也是 DLight 能跨 CUDA/ROCm/Vulkan/Metal 通用的一条 GPU 规则的原因。8.3benchmark调度结果的量化评估python/tvm/s_tir/dlight/benchmark 提供三个模块bench.py基准测试入口对给定调度后的模块执行计时extract.py从调度结果中提取可对比的统计信息utils.py共用工具。如果你在自定义规则强烈建议用benchmark验证自己的规则相对Fallback或手工 baseline 的实际收益避免看起来漂亮、跑起来没差的规则。九、编写自定义 DLight 规则的完整流程综合以上分析落地一条自定义规则的推荐路径是分析用dl.SBlockInfo/IterInfo或直接遍历 TIR 判断函数是否符合你的算子模式实现继承ScheduleRule实现apply或直接dl.ScheduleRule.from_callable(MyRule)包装普通函数参数化用dl.max_threads_per_block(target)、dl.suggest_threads_per_block(...)、dl.auto_vectorize(...)等工具按 target 推导配置不要写死常量兜底在你的规则链末尾放FallbackGPU/Adreno以保证覆盖率验证用dl.benchmark量化收益用ApplyDefaultSchedule把规则挂进自己的 pass 管线参考 Relax CUDA 后端的写法。示例骨架from tvm.s_tir import dlight as dl from tvm import tirx from tvm.target import Target dl.ScheduleRule.from_callable(MyElemwiseRule) def my_elemwise(func: tirx.PrimFunc, target: Target, tunable: bool): # 1. 结构判断不是 elementwise 就返回 None if not _is_elemwise(func): return None # 2. 创建 schedule 并调度 sch tirx.create_schedule(func) block ... # 定位目标块 sch.parallel(block) # 或 bind/vectorize sch.vectorize(block) # 配合 auto_vectorize return sch # 接入自己的编译管线 pipeline_passes.append( dl.ApplyDefaultSchedule(my_elemwise, dl.gpu.Fallback()) )十、总结DLight 的定位与适用边界维度结论定位为常见深度学习负载提供零调优的默认调度同时可为 MetaSchedule 提供可调优空间tunableTrue入口dl.ApplyDefaultSchedule(rule1, rule2, ...)模块 pass规则列表按专用在前、兜底在后排序规则组织顶层聚合导出gpu/cpu/adreno按硬件划分analysis/base/benchmark提供支撑集成方式已内置于 Relax 各后端默认管线见 cuda/pipeline.py也可独立使用扩展方式实现ScheduleRule.apply或用from_callable包装函数配合analysis与base.utils工具适用边界覆盖 GEMM/GEMV/Reduction/Transpose/RMSNorm/Conv/Pool 等主流模式不在覆盖范围内的算子依赖Fallback兜底追求极致性能时可退化为手工 schedule 或 MetaSchedule 调优DLight 的本质是把社区积累的调度经验沉淀为可组合、可扩展、按 target 自适应的规则库。理解它的 API 分层gpu/cpu/adreno规则集 analysis分析层 base基础设施 benchmark验证层之后无论你是想快速部署模型、为特定算子定制调度还是想为 MetaSchedule 扩充搜索空间都能在tvm.s_tir.dlight找到对应的抓手。更多细节可继续阅读 docs/reference/api/python/s_tir/dlight.rst 对应的 automodule 文档以及 python/tvm/s_tir/dlight 下的源码注释。【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址: https://gitcode.com/gh_mirrors/tv/tvm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 TVM s_tir.dlight 模块完全指南：开箱即用的默认调度规则体系与 GPU/CPU/Adreno 实战
TVM s_tir.dlight 模块完全指南开箱即用的默认调度规则体系与 GPU/CPU/Adreno 实战【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址: https://gitcode.com/gh_mirrors/tv/tvm导读tvm.s_tir.dlight是 TVM 中面向深度学习负载的开箱即用调度库它内置了一批针对 GEMM、GEMV、Reduction、Transpose、RMSNorm 等典型算子的调度规则ScheduleRule无需任何调优即可为 IRModule 中的每个PrimFunc生成高质量的默认调度也可以为 MetaSchedule 提供可调优的搜索空间。本文以 docs/reference/api/python/s_tir/dlight.rst 为骨架结合仓库中 python/tvm/s_tir/dlight 的源码实现与 Relax 后端集成路径系统讲解该模块的公共 API、各子模块职责、调度规则的运行原理以及如何通过ApplyDefaultSchedule接入自定义编译管线。一、模块定位DLight 是什么解决什么问题DLightDeep Learning Lightweight Scheduler的核心目标正如其包级 docstring 所言provide efficient schedules out-of-box for deep learning workloadspython/tvm/s_tir/dlight/init.py。它在传统两条路线之间取了一个平衡点完全手工写 schedule需要开发者对每个算子和每种硬件写大量T.schedule指令成本高、难维护MetaSchedule 自动调优质量上限高但需要 builder/runner 实测、代价大不适合推理部署这类要求快速出结果的场景。DLight 的做法是把常见算子模式的调度经验固化成规则rule在编译期对 TIR 层面做结构分析直接命中规则并产出完整 schedule从而做到零调优、开箱即用。从文档的 RST 结构看tvm.s_tir.dlight的公开面被组织成七个部分docs/reference/api/python/s_tir/dlight.rstRST 章节对应包路径职责tvm.s_tir.dlightpython/tvm/s_tir/dlight/init.py顶层统一入口聚合导出规则类、分析工具与基础设施tvm.s_tir.dlight.gpupython/tvm/s_tir/dlight/gpu/init.pyGPU 通用调度规则CUDA/ROCm/Vulkan/Metal 均适用tvm.s_tir.dlight.adrenopython/tvm/s_tir/dlight/adreno/init.pyAdreno 移动 GPU 专用规则卷积/池化/布局转换tvm.s_tir.dlight.cpupython/tvm/s_tir/dlight/cpu/init.pyCPU 通用规则GEMV、Reductiontvm.s_tir.dlight.analysispython/tvm/s_tir/dlight/analysis调度前结构分析工具SBlockInfo、IterInfo 等tvm.s_tir.dlight.basepython/tvm/s_tir/dlight/base基础设施ScheduleRule 抽象、ApplyDefaultSchedule pass、公共工具函数tvm.s_tir.dlight.benchmarkpython/tvm/s_tir/dlight/benchmark调度结果提取与基准测试辅助注gpu子模块在 docstring 中特别说明For CUDA/ROCm/Vulkan/Metal-specific rules, usetvm.s_tir.dlight.cuda/rocm/vulkan/metalinstead即这里的规则是GPU 通用而非某厂商私有的。二、顶层公共 APItvm.s_tir.dlight导出内容顶层__init__.pypython/tvm/s_tir/dlight/init.py共导出三类对象后端子模块gpu、adreno、cpu分别承载对应硬件的规则集合分析工具来自dlight.analysisSBlockInfo对 PrimFunc 的调度块SBlock做结构信息封装IterInfo描述单个循环迭代变量的信息是否归约、是否空间等normalize_prim_func把函数归一化为便于规则匹配的标准形态基础设施来自dlight.baseApplyDefaultSchedule一个module_pass的 IRModule pass把一组规则应用到模块内所有 PrimFunc 上ScheduleRule规则的抽象基类try_inline/try_inline_contiguous_spatial通用的 producer/consumer 内联辅助函数。因此用户最常见的导入方式就是from tvm.s_tir import dlight as dl # 使用 GPU 通用规则 dl.gpu.Matmul() # 使用 CPU 规则 dl.cpu.Reduction() # 使用分析工具 dl.SBlockInfo(func) # 使用默认调度 pass dl.ApplyDefaultSchedule(dl.gpu.Matmul(), dl.gpu.Fallback())三、ScheduleRule抽象规则的定义与生命周期ScheduleRule定义在 python/tvm/s_tir/dlight/base/schedule_rule.py是整个 DLight 体系最核心的接口。它的语义是A thin wrapper on an arbitrary function that can be used to schedule a TIR PrimFunc.3.1apply方法核心契约每个规则必须实现apply(func, target, tunable)functirx.PrimFunc待调度的 TIR 函数targettvm.target.Target规则为其生成 schedule 的编译目标tunable布尔标志控制规则产出的形式tunableFalse只能返回一个s_tir.Schedule且其中不允许包含可调优指令tunable instructions——这是默认调度模式的约束tunableTrue可以返回一个或一组s_tir.Schedule且允许包含可调优指令——这是为 MetaSchedule 提供搜索空间的模式返回值可以是None、单个s_tir.Schedule或list[s_tir.Schedule]。返回None表示当前规则不适用于该 PrimFunc调度引擎会继续尝试下一条规则。3.2from_callable用普通函数快速声明规则ScheduleRule.from_callable(name)是一个装饰器工厂让你不必继承基类就能把一个普通函数包装成规则from tvm.s_tir import dlight as dl from tvm import tirx from tvm.target import Target dl.ScheduleRule.from_callable(MyRule) def my_rule(func: tirx.PrimFunc, target: Target, tunable: bool): # 对 func 做结构判断返回 None 或 s_tir.Schedule ...其内部实现是为传入函数动态生成一个_Rule子类并把函数绑定到apply上同时将类名设置为传入的name源码见 schedule_rule.py。3.3is_target_available默认返回True子类可重写用于在特定 target如无 CUDA 的机器上提前声明规则不可用从而被调度引擎跳过。3.4 规则匹配的调度语义以_apply_rules为例ApplyDefaultSchedule内部通过_apply_rulestransform.py串行遍历规则列表for rule in rules: space rule.apply(func, target, tunable) if space is None: continue # 当前规则不适用尝试下一条 if isinstance(space, s_tir.Schedule): space [space] return space # 命中第一条规则即返回 return None这种顺序尝试、先命中者胜出的机制决定了规则列表的排列顺序很有讲究越通用的兜底规则如Fallback应排在最后越专门的规则如Matmul排在前面。这一点可以从 Relax 的 CUDA 后端默认管线的排列得到印证见下文第六节。四、ApplyDefaultSchedule把规则变成一次 IRModule passApplyDefaultSchedule定义在 python/tvm/s_tir/dlight/base/transform.py是用户接触最多的入口。它是一个module_pass(opt_level0, nameApplyDefaultSchedule)即属于module pass输入输出都是IRModuleopt_level0意味着它在 pass 管线中属于功能性必做而非优化可选类别构造时传入若干ScheduleRule实例ApplyDefaultSchedule(*rules)。4.1 工作流程transform_module对模块内每个函数执行如下逻辑transform.py只处理tirx.PrimFunc且跳过已调度过的函数——判断依据是函数属性tirx.is_scheduled是否为 1_is_scheduled获取目标设备优先读取函数自带的target属性否则回退到Target.current()_get_targetallow_noneFalse表示当前无 target 上下文会直接报错以tunableFalse调用_apply_rules若命中规则将产出的 schedule 对应的新函数写回模块并打上tirx.is_scheduled True标记避免后续被重复调度。4.2 幂等性与去重tirx.is_scheduled标记是这套机制的关键设计它保证了 pass 可重复执行、可与其它 pass 任意组合而不会产生二次调度。自定义编译器如果自己写了调度逻辑同样可以依赖该属性避免与 DLight 冲突。4.3 直接用法示例import tvm from tvm import tirx from tvm.s_tir import dlight as dl from tvm.script import tirx as T T.prim_func def my_gemm(A: T.Buffer((256, 256), float32), B: T.Buffer((256, 256), float32), C: T.Buffer((256, 256), float32)): for i, j, k in T.grid(256, 256, 256): with T.block(C): vi, vj, vk T.axis.remap(SSR, [i, j, k]) with T.init(): C[vi, vj] 0.0 C[vi, vj] A[vi, vk] * B[vk, vj] mod tvm.IRModule({main: my_gemm}) target tvm.target.Target(cuda) # 应用 DLight 默认调度 with target: mod dl.ApplyDefaultSchedule(dl.gpu.Matmul(), dl.gpu.Fallback())(mod) # 之后即可正常走 codegen 生成 CUDA 内核说明上述tunableFalse的默认路径对每个函数只生成一个确定 schedule若想拿到一组候选空间用于自动调优可参考ScheduleRule.apply(tunableTrue)的语义结合 MetaSchedule 的搜索策略使用。五、GPU 规则集tvm.s_tir.dlight.gpugpu子模块python/tvm/s_tir/dlight/gpu/init.py是覆盖度最广的规则集合共导出 8 条规则规则类源文件适用算子模式Matmulgpu/matmul.pyGEMM含 Batch GEMM如全连接、attention 的 QK^T / PVGEMVgpu/gemv.py矩阵-向量乘M 很小如解码阶段的 LinearLowBatchGEMVgpu/low_batch_gemv.pybatch 维度很小的 GEMV进一步特化Reductiongpu/reduction.py规约类算子sum/max/min/softmax 的一部分等GeneralReductiongpu/general_reduction.py一般化规约模式的兜底Transposegpu/transpose.py布局转置算子RMSNormgpu/rmsnorm.pyRMSNorm / LayerNorm 类归一化Fallbackgpu/fallback.py兜底规则上述规则都不命中时保证仍有可用 schedule5.1 源码级理解以Matmul为例Matmul规则gpu/matmul.py是理解整条规则链的绝佳样例其调度流程体现了 DLight 的通用套路结构分析调用get_root_block、IterInfo/SBlockInfo等分析工具识别 GEMM 的三个迭代维度M、N、K区分空间轴与归约轴内联辅助算子利用auto_inline_producers/auto_inline_consumers反复对 producer/consumer 块执行compute_inline直到没有可内联的块为止见 matmul.py把 GEMM 周边 elementwise 算子全部折叠进主计算分块与线程绑定对 M、N 维度做 tile并基于 target 信息做 block/thread 绑定与向量化返回 Schedule生成完整s_tir.Schedule供后续 codegen。auto_inline_*的实现非常直接递归收集所有 producer/consumer 块逐个尝试sch.compute_inline(block)抛异常不可内联就跳过继续下一个直到一轮下来没有任何内联成功才停止。这种能内联就内联的激进策略是 DLight 保持内核简洁高效的基石。5.2Fallback规则的意义Fallback保证任何函数都能被调度因此它是所有 GPU 规则链的最后一道保险。在自定义编译管线中如果你不确定自己的算子集合能被 DLight 覆盖务必把Fallback放在规则列表末尾避免_apply_rules返回None导致整个 pass 对该函数无所作为。六、真实接入案例Relax CUDA 后端如何内置 DLightDLight 并不是一个孤立的工具库它已经被整合进 TVM 的 Relax 编译流程。以 CUDA 后端为例python/tvm/relax/backend/cuda/pipeline.py 中的legalize_passes展示了标准接入姿势def legalize_passes(target: tvm.target.Target): from tvm.s_tir import dlight as dl return [ tvm.relax.transform.LegalizeOps(), tvm.relax.transform.AnnotateTIROpPattern(), tvm.relax.transform.FoldConstant(), tvm.relax.transform.FuseOps(), tvm.relax.transform.FuseTIR(), # 先完成算子融合产出 TIR PrimFunc dl.ApplyDefaultSchedule( # 再对融合后的函数应用 DLight 默认调度 dl.gpu.Matmul(), dl.gpu.GEMV(), dl.gpu.Reduction(), dl.gpu.GeneralReduction(), dl.gpu.Fallback(), ), ]这串代码揭示了 DLight 在完整编译流程中的精确位置FuseOpsFuseTIR先把 Relax 层的算子图融合成一个个 TIRPrimFuncApplyDefaultSchedule紧接着对这些融合函数按Matmul → GEMV → Reduction → GeneralReduction → Fallback的顺序尝试匹配规则之后才进入 dataflow lowering 与最终 codegen见同一文件中的dataflow_lower_passes与finalize_passes。也就是说DLight 取代了传统手工写调度如sch tirx.create_schedule后逐条 bind/vectorize的环节用户只需声明规则列表编译管线自动完成调度。类似的集成还出现在 python/tvm/relax/backend/rocm/pipeline.py、python/tvm/relax/backend/gpu_generic/pipeline.py、python/tvm/relax/backend/adreno/pipeline.py 以及 python/tvm/relax/pipeline.py 中。七、CPU 与 Adreno 规则集7.1 CPUtvm.s_tir.dlight.cpuCPU 规则集python/tvm/s_tir/dlight/cpu/init.py相对精简导出两条CPU 通用规则GEMVcpu/gemv.py矩阵-向量乘重点处理缓存友好的分块与向量化Reductioncpu/reduction.py规约类算子关注寄存器/多线程规约的组织。CPU 场景下 DLight 的核心收益是快速生成多线程 向量化的内核无需手工做parallel/vectorize决策。7.2 Adrenotvm.s_tir.dlight.adrenoAdreno高通移动 GPU规则集python/tvm/s_tir/dlight/adreno/init.py面向移动端推理导出四条规则Conv2dadreno/convolution.py二维卷积移动端最核心的算子Pool2Dadreno/pool.py二维池化LayoutTransformadreno/layout_transform.py布局转换如 NHWC 相关重排Fallbackadreno/fallback.py兜底规则。Adreno 的调度重点关注移动 GPU 的纹理/局部内存特性结合 python/tvm/relax/backend/adreno/pipeline.py 可以看到它在 Relax Adreno 后端管线中同样承担默认调度职责。八、分析工具与基础设施analysis、base、benchmark8.1analysis调度前的结构分析python/tvm/s_tir/dlight/analysis 提供规则共享的分析设施SBlockInfo、IterInfo把 TIR 块与循环迭代结构抽象成规则易于查询的数据结构顶层__init__.py直接再导出见 dlight/init.pycommon_analysis.py常见分析的共享实现gemv.py面向 GEMV 模式的结构识别。规则编写者应尽量复用这些分析工具而不是各自解析 TIR 语法树从而保证整个规则库对结构判断的口径一致。8.2base基础设施与工具函数base包python/tvm/s_tir/dlight/base除了前文讲过的ScheduleRule、ApplyDefaultSchedule、try_inline/try_inline_contiguous_spatial定义于 base/common_schedules.py还导出一组调度参数启发式工具base/utils.pymax_threads_per_block查询目标允许的最大线程数/块大小供规则决定 block 配置suggest_threads_per_block根据循环长度启发式地给出线程数建议auto_vectorize自动选择向量化宽度get_extent/get_bytes读取循环 extent 与数据类型字节数用于计算分块大小。这些函数是 DLight 规则按 target 自适应的秘密武器规则不写死线程数而是借助这些查询函数在编译期动态推导这也是 DLight 能跨 CUDA/ROCm/Vulkan/Metal 通用的一条 GPU 规则的原因。8.3benchmark调度结果的量化评估python/tvm/s_tir/dlight/benchmark 提供三个模块bench.py基准测试入口对给定调度后的模块执行计时extract.py从调度结果中提取可对比的统计信息utils.py共用工具。如果你在自定义规则强烈建议用benchmark验证自己的规则相对Fallback或手工 baseline 的实际收益避免看起来漂亮、跑起来没差的规则。九、编写自定义 DLight 规则的完整流程综合以上分析落地一条自定义规则的推荐路径是分析用dl.SBlockInfo/IterInfo或直接遍历 TIR 判断函数是否符合你的算子模式实现继承ScheduleRule实现apply或直接dl.ScheduleRule.from_callable(MyRule)包装普通函数参数化用dl.max_threads_per_block(target)、dl.suggest_threads_per_block(...)、dl.auto_vectorize(...)等工具按 target 推导配置不要写死常量兜底在你的规则链末尾放FallbackGPU/Adreno以保证覆盖率验证用dl.benchmark量化收益用ApplyDefaultSchedule把规则挂进自己的 pass 管线参考 Relax CUDA 后端的写法。示例骨架from tvm.s_tir import dlight as dl from tvm import tirx from tvm.target import Target dl.ScheduleRule.from_callable(MyElemwiseRule) def my_elemwise(func: tirx.PrimFunc, target: Target, tunable: bool): # 1. 结构判断不是 elementwise 就返回 None if not _is_elemwise(func): return None # 2. 创建 schedule 并调度 sch tirx.create_schedule(func) block ... # 定位目标块 sch.parallel(block) # 或 bind/vectorize sch.vectorize(block) # 配合 auto_vectorize return sch # 接入自己的编译管线 pipeline_passes.append( dl.ApplyDefaultSchedule(my_elemwise, dl.gpu.Fallback()) )十、总结DLight 的定位与适用边界维度结论定位为常见深度学习负载提供零调优的默认调度同时可为 MetaSchedule 提供可调优空间tunableTrue入口dl.ApplyDefaultSchedule(rule1, rule2, ...)模块 pass规则列表按专用在前、兜底在后排序规则组织顶层聚合导出gpu/cpu/adreno按硬件划分analysis/base/benchmark提供支撑集成方式已内置于 Relax 各后端默认管线见 cuda/pipeline.py也可独立使用扩展方式实现ScheduleRule.apply或用from_callable包装函数配合analysis与base.utils工具适用边界覆盖 GEMM/GEMV/Reduction/Transpose/RMSNorm/Conv/Pool 等主流模式不在覆盖范围内的算子依赖Fallback兜底追求极致性能时可退化为手工 schedule 或 MetaSchedule 调优DLight 的本质是把社区积累的调度经验沉淀为可组合、可扩展、按 target 自适应的规则库。理解它的 API 分层gpu/cpu/adreno规则集 analysis分析层 base基础设施 benchmark验证层之后无论你是想快速部署模型、为特定算子定制调度还是想为 MetaSchedule 扩充搜索空间都能在tvm.s_tir.dlight找到对应的抓手。更多细节可继续阅读 docs/reference/api/python/s_tir/dlight.rst 对应的 automodule 文档以及 python/tvm/s_tir/dlight 下的源码注释。【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址: https://gitcode.com/gh_mirrors/tv/tvm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考