ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

TVM s_tir.schedule 调度 API 完全指南:TensorIR 可调度编程模型与调度原语详解

TVM s_tir.schedule 调度 API 完全指南:TensorIR 可调度编程模型与调度原语详解 模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载导读tvm.s_tir.schedule是 TVM 开源编译器中面向 TensorIRscheduable TensorIR即 S-TIR的调度ScheduleAPI 命名空间。本文以该模块的官方 API 文档docs/reference/api/python/s_tir/schedule.rst为骨架结合模块源码 python/tvm/s_tir/schedule/schedule.py 与配套的 state.py、transform.py 等实现文件系统讲解 Schedule 类、随机变量RV、ScheduleState、Trace 等核心概念并逐一剖析拆分、融合、重排、绑定线程、缓存、compute-at、张量化、布局变换、填充分解等数十个调度原语的语义、前置条件与代码示例。读完本文你将掌握如何用 Python 脚本编写一套完整的 TensorIR 调度流程并能理解其底层 sref 树与依赖分析的运行机制。一、模块定位从 RST 引用页到完整调度命名空间docs/reference/api/python/s_tir/schedule.rst是一个典型的 Sphinxautomodule引用页其正文由以下指令生成tvm.s_tir.schedule ------------------ .. automodule:: tvm.s_tir.schedule :members: :imported-members:也就是说该页面会自动收集tvm.s_tir.schedule模块中所有公开成员及其 docstring构成完整的 API 文档。文档的实质内容全部沉淀在源码 docstring 中因此本指南直接以源码中的类与方法定义为准。从 python/tvm/s_tir/schedule/init.py 可以看到该命名空间的导出面from ..sblock_scope import SBlockScope, Dependency, DepKind, StmtSRef from .instruction import Instruction, InstructionKind from .schedule import SBlockRV, ExprRV, LoopRV, Schedule, ScheduleError from .state import ScheduleDebugMask, ScheduleState from .trace import Trace from . import analysis from . import transform模块内部文件分工如下文件职责schedule.py用户面对的核心Schedule类以及LoopRV、SBlockRV、ExprRV三种随机变量state.pyScheduleState调度的核心数据结构与ScheduleDebugMask调试位掩码trace.pyTrace记录调度指令执行轨迹支撑可复现与决策回放instruction.pyInstruction/InstructionKind单条调度指令的表示transform.pytile_with_tensor_intrin等高级组合变换analysis.py供调度使用的分析工具testing.py调度相关的测试辅助sblock_scope.pySBlockScope、Dependency、DepKind、StmtSRef依赖分析设施二、核心概念Schedule、随机变量与 ScheduleState2.1 Schedule 类语义保持的变换集合Schedule是用户直接使用的调度类源码 docstring 对其定义如下schedule.pyA schedule is a set of transformations that change the order of computation but preserve the semantics of computation. Some example of schedules:Split a loop into two;Reorder two loops;Inline the computation of a specific buffer into its consumer即调度是一组改变计算顺序但不改变计算语义的变换典型例子包括把一个循环拆成两个、重排两个循环、将某个 buffer 的计算内联进它的消费者。2.2 三种随机变量Random Variable调度过程不直接操作 IR 节点而是操作随机变量由Schedule内部的符号表维护其求值关系schedule.pyLoopRV指代一个循环For节点的随机变量SBlockRV指代一个 block的随机变量ExprRV指代一个整数表达式的随机变量ExprRV Expr的类型别名。通过sch.get(rv)可以把随机变量求值为对应的SBlock/For/ 整数通过sch.get_sref(stmt_or_rv)则可以拿到指向 IR 语句的StmtSRef引用。2.3 ScheduleState调度的核心数据结构ScheduleStatestate.py是调度原语真正操作的底层状态包含被调度的 ASTmod可调度语句的sref 树由StmtSRef构成每个 block 作用域的依赖信息block_info从 AST 节点到 sref 树节点的反向映射调试标志debug_mask开启后做额外正确性检查前置检查开关enable_check。CachedFlags namedtuple(CachedFlags, [affine_binding, region_cover, stage_pipeline])缓存了每个 block 的三个关键属性调度原语的大量前置条件都依赖它们。2.4 调试位掩码ScheduleDebugMaskstate.py 定义了ScheduleDebugMask枚举VERIFY_SREF_TREE 1每条调度指令后验证 sref 树正确性VERIFY_CACHED_FLAGS 2验证affine_binding、region_cover、stage_pipeline缓存标志。构造 Schedule 时debug_maskall等价于打开上述两位debug_masknone等价于0也可直接传整数位掩码。三、Schedule 的构造与参数详解构造Schedule只需传入一个PrimFunc或IRModuleschedule.pysch tvm.s_tir.Schedule( mod, # PrimFunc | IRModule seedNone, # 随机种子 debug_masknone, # all / none / int 位掩码 error_render_leveldetail, # detail / fast / none enable_checkTrue, # 是否执行前置条件检查 )各参数语义与边界条件mod接受PrimFunc或IRModule。传入PrimFunc时会被自动包装为IRModule({main: func})见 state.py 的_parse_mod。seedNone与-1表示使用设备随机源否则必须为[1, 2147483647]内的整数schedule.py。debug_mask字符串 all/none 或整数位掩码控制创建类之后及每次Replace调用之后是否做额外正确性检查检查项包括 VerifySRefTree 与 VerifyCachedFlags。error_render_level错误渲染级别。detail 输出带 TIR 与错误位置的详细信息fast 只输出简单错误信息不做渲染与字符串处理none 不显示任何错误信息。enable_check默认调度前置检查较严格可能阻碍某些合法调度。设为False可跳过部分前置检查但正确性由用户自己保证——即使跳过检查调度真正失败时依然会抛错。构造时底层调用的是_ffi_api.TracedSchedule带轨迹记录Schedule._create_non_traced则通过_ffi_api.ConcreteSchedule构造不带轨迹的调度器供无需回放决策的高性能场景使用。四、常用工具方法状态访问、工作函数切换与复制Schedule提供了一批便捷属性和方法schedule.py方法说明sch.mod返回当前被调度的 IRModule ASTsch.state返回当前ScheduleStatesch.trace返回内部维护的调度指令轨迹Trace非 traced 调度返回Nonesch.func_working_on返回当前正在调度的函数的GlobalVarsch.work_on(func_name)显式指定要调度的函数默认调度名为main的函数或 IRModule 中唯一函数多函数且无 main 时必须调用它sch.copy()深拷贝调度sref 树完全重建、被调度 IRModule 保持不变、所有随机变量在新拷贝中依然有效sch.seed(seed)重新播种随机数-1表示设备随机源sch.fork_seed()返回一个分叉的新随机状态用于派生新调度sch.show()打印高亮 TVM script同时打印模块与 trace参数转发给Module.show与Trace.show五、查找与随机变量求值sch.get(rand_var_or_sref)将ExprRV/SBlockRV/LoopRV/StmtSRef求值为对应整数、SBlock或Forschedule.py。求值结果为IntImm时自动转为 Pythonint。sch.get_sref(rand_var_or_stmt)返回对应StmtSRefschedule.py。sch.remove_rv(rand_var)从符号表移除一个随机变量schedule.py。六、随机采样Sampling为自动调优提供决策空间自动调度如 meta-schedule需要随机采样参数Schedule提供四类采样原语schedule.pysample_categorical(candidates, probs, decisionNone) - ExprRV按概率分布probs从候选整数candidates中采样sample_perfect_tile(loop, n, max_innermost_factor16, decisionNone) - list[ExprRV]对指定循环采样n个完美分块因子innermost 因子不超过max_innermost_factorsample_partitioned_tile(loop, n, partition_pos0, innerpart_factor1, decisionNone) - list[ExprRV]采样分区partitioned分块因子partition_pos指定将 tiles 分成两段的位置innerpart_factor为第二段因子sample_compute_location(block, decisionNone) - LoopRV为给定 block 采样 compute-at 位置。每个采样原语都可传decision显式给出决策值配合Trace实现完全可复现的调度决策。七、获取 Block 与循环调度第一步通常是定位要变换的 block 和循环schedule.pyget_sblock(name, func_nameNone) - SBlockRV按名字取 block。默认在work_on指定的函数内查找若名字匹配 0 个或多个 block 则抛IndexErrorget_loops(block) - list[LoopRV]返回 block 作用域内的父循环从外到内排序get_child_blocks(block_or_loop) - list[SBlockRV]返回某 block/loop 内的叶子 blockget_producers(block)/get_consumers(block)分别返回指定 block 的生产者 / 消费者 block 列表依赖分析结果基于SBlockScopeget_output_blocks(scope_block) - list[SBlockRV]返回作用域内所有输出 block——即写入了某些 buffer 但该 buffer 并非在 PrimFunc 内分配的 block。八、循环变换fuse / split / merge / reorder / partition8.1 fuse合并相邻循环fuse(*loops, preserve_unit_itersTrue) - LoopRVschedule.py将一串连续循环融合为一层。前置条件循环不能带 annotation 或 thread binding第i1个循环必须是第i个循环的唯一子节点所有循环必须从 0 开始被融合循环的 domain 不能依赖于其他待融合循环。T.prim_func(s_tirTrue) def before_fuse(a: T.handle, b: T.handle) - None: A T.match_buffer(a, (128, 128)) B T.match_buffer(b, (128, 128)) for i, j in T.grid(128, 128): with T.sblock(B): vi, vj T.axis.remap(SS, [i, j]) B[vi, vj] A[vi, vj] * 2.0 sch tvm.s_tir.Schedule(before_fuse) i, j sch.get_loops(sch.get_sblock(B)) sch.fuse(i, j)融合后两层128循环变成一层0..16384的循环for i_j_fused in T.serial(0, 16384): with T.sblock(B): vi T.axis.S(128, T.floordiv(i_j_fused, 128)) vj T.axis.S(128, T.floormod(i_j_fused, 128)) B[vi, vj] A[vi, vj] * 2.08.2 split拆分循环split(loop, factors, preserve_unit_itersTrue, disable_predicationFalse) - list[LoopRV]schedule.py把循环按因子列表拆成连续多层。factors中最多允许一个None自动推断补齐也支持ExprRV与正整数字面量必要时会插入 predicate 保证迭代总数不变。sch tvm.s_tir.Schedule(before_split) i, j sch.get_loops(sch.get_sblock(B)) sch.split(i, factors[2, 64])拆分后for i0, i1, j in T.grid(2, 64, 128)且 block 绑定相应更新为vi T.axis.S(128, i0 * 64 i1)。参数disable_predicationTrue适用于可扩展scalable因子且调度作者已知因子能整除循环界的场景可避免生成 guard 谓词但误用可能导致错误的代码生成源码以 Warning 明确标注。8.3 merge合并同层循环merge(*loops) - LoopRVschedule.py合并两个原本同层并列的循环如两个独立的计算阶段各自的 i 循环要求它们处于同一作用域、无 annotation/thread binding、从 0 开始且 extent 与嵌套深度相同且从目标循环到其 LCA 之间内层循环必须是外层唯一子节点。注意fuse作用于嵌套循环merge作用于并列循环二者语义不同。8.4 reorder 与 reorder_block_iter_varreorder(*ordered_loops)schedule.py按新顺序重排同一链条上的循环要求循环处于同一链条可排成[l_1, l_2, ..., l_n]l_i是l_{i1}的祖先中间只有单分支循环重排后外层循环的 domain 不能依赖任何内层循环所有 block 绑定必须是仿射的block 变量只能是数据并行或归约参数中不允许重复循环。sch.reorder(j, i) # 把 (i, j) 变成 (j, i)reorder_block_iter_var(block, new_order)schedule.py则直接重排 block 内部 itervar 的顺序不改变外层循环适用于 matmul 这类 block 绑定为 SSR 的情形sch.reorder_block_iter_var(C, [2, 1, 0]) # 将 [vi, vj, vk] 变为 [vk, vj, vi]8.5 loop_partition循环分区loop_partition(loop, factors, preserve_unit_itersTrue) - list[LoopRV]schedule.py按因子把循环分成多个连续区段不同于 split 的嵌套拆分每个区段拥有独立循环与各自的访问区域factors同样最多允许一个None。典型输出会把原始循环拆为 common 段 若干 partition 段 的嵌套结构。8.6 add_unit_loop添加单位循环add_unit_loop(block_or_loop) - LoopRVschedule.py在指定 block 或循环之上包一层 extent 为 1 的单位循环for u in T.serial(1)为后续 bind/parallel 等提供新的可调度循环层级。九、控制并行度parallel / vectorize / bind / unroll这四个原语修改循环的ForKindschedule.pyparallel(loop)把循环标记为T.parallel。要求所在作用域 block 具备 stage-pipeline 属性、循环下所有 block 是 complete/reduction block 且绑定仿射循环只能出现在数据并行 itervar 的绑定中vectorize(loop)把循环标记为T.vectorized前置条件与parallel相同bind(loop, thread_axis)把循环绑定到线程轴thread_axis可选blockIdx.x/y/z、threadIdx.x/y/z、vthread.x/y/z裸vthread为将被弃用的遗留行为。threadIdx.*可出现在数据并行与归约 itervar 绑定中其余线程轴只能出现在数据并行绑定中sch.bind(i, blockIdx.x) sch.bind(j, threadIdx.x)绑定后 IR 变为for i in T.thread_binding(0, 128, threadblockIdx.x): for j in T.thread_binding(0, 128, threadthreadIdx.x): with T.sblock(B): vi, vj T.axis.remap(SS, [i, j]) B[vi, vj] A[vi, vj] * 2.0unroll(loop)把循环标记为T.unroll无额外前置条件It requires nothing。十、缓存阶段插入cache_read / cache_write 及其变体缓存类原语用于在存储层级如 local / shared中插入中间缓存 stage是访存优化的核心手段。10.1 cache_read 与 cache_writecache_read(block, read_buffer_index, storage_scope, consumer_blocksNone) - SBlockRVschedule.py为消费者 block 的某个读 buffer 生成读缓存 stage。前置条件作用域内写该 buffer 的 block 至多一个且作用域 block 具备 stage-pipeline 属性。sch.cache_read(block_b, 0, local)生成一个A_local T.sblock_alloc_buffer((128, 128), scopelocal)的缓存 block原 block 的读全部改走A_local。cache_write(block, write_buffer_index, storage_scope, consumer_blocksNone) - SBlockRVschedule.py则相反为生产者 block 的写 buffer 生成写缓存写缓存先接收计算结果再由原 block 从缓存拷回。两者的 buffer 参数都支持三种形式整数索引读/写区域下标、唯一 buffer 名字字符串、Buffer对象consumer_blocks可指定哪些消费者走缓存不指定则全部。10.2 reindex_cache_read / reindex_cache_write这两个变体schedule.py用自定义index_maplambda或IndexMap指定缓存 buffer 的排布从而在缓存的同时完成索引重排如转置(vj, vi)或折叠(vi // 2, vi % 2, vj)。它们只支持单一消费者多消费者请使用cache_read/cache_write缓存 block 会省略未出现在 buffer 访问区中的 itervar可用transform_block_layout继续调整。10.3 cache_inplace 与 cache_indexcache_inplace(block, read_buffer_index, storage_scope) - list[SBlockRV]schedule.py针对同时读写同一 bufferinplace 操作如T.call_extern就地修改的 block在操作前拷入缓存、操作后拷回返回[读缓存 block, 写缓存 block]cache_index(block, storage_scope, cse_thresh0) - list[SBlockRV]schedule.py把 block 中重复计算的下标表达式预计算到缓存 buffer如 resize 中的vi//4 vj//4cse_thresh为公共子表达式判定阈值默认 0 表示缓存全部下标计算。10.4 reindexreindex(block, buffer) - SBlockRVschedule.py为访问某 buffer 的 block 生成按 block 迭代器排布的新缓存要求只有一个 block访问该 buffer、且 block 内该 buffer 的 load/store 各只有一处。buffer参数支持(read/write, index)元组、buffer 名字或Buffer对象。十一、计算位置调整compute_at 家族11.1 compute_atcompute_at(block, loop, preserve_unit_loopsFalse, index-1)schedule.py把生产者 block 移动到指定循环之下并重新生成循环使生产区域恰好覆盖该循环下所有消费者的消费区域。前置条件block与loop在同一作用域且loop不是block的祖先作用域 block 具备 stage-pipeline 属性作用域子树满足紧凑数据流条件所有 block 是 complete 或 reduction blockblock相对作用域不是输出 block其写出的 buffer 在作用域内分配block的所有消费者都在给定loop之下。index控制插入位置-1为最后一个可行插入点-2为第一个可行插入点非负整数为显式插入点。block sch.get_sblock(B) loop, _ sch.get_loops(sch.get_sblock(C)) sch.compute_at(block, loop, preserve_unit_loopsFalse)reverse_compute_atschedule.py方向相反把消费者 block移动到某循环下使消费区域覆盖该循环下所有生产者的产出区域其余条件对称。11.2 compute_inline 与 reverse_compute_inlinecompute_inline(block)schedule.py把 block 内联到其消费者中消除中间 buffer# beforeB[vi, vj] A[vi, vj] * 2.0C[vi, vj] B[vi, vj] 1.0 sch.compute_inline(sch.get_sblock(B)) # afterC[vi, vj] A[vi, vj] * 2.0 1.0B 及其中间 buffer 消失前置条件block 是只产出一个 buffer 的 complete 非 root blockblock 不能是作用域内唯一叶子block 体必须是形如A[i, j, k, ...] ...的BufferStoreLHS 下标全是互不相同的原子变量。reverse_compute_inline(block)schedule.py把 block 内联进其唯一生产者要求该 block 只生产且只消费一个 buffer且生产者是读后写RAW的 complete 非 root block。11.3 fuse_reduction_epiloguefuse_reduction_epilogue(reduction_block, epilogue_block)schedule.py把 epilogue尾算子融合进归约 block支持三类模式Biasoutput reduction_result bias、BiasReLUoutput max(reduction_result bias, 0)、Clippingoutput min(max(reduction_result, lower), upper)。源码对此给出重要语义警告BiasReLU 与 Clipping 这类非线性epilogue 融合后语义从归约后应用变为逐迭代应用数值结果可能不同如对[6, -2]做clip([-5,5])原始clip(sum)clip(4)4融合后acc0 → clip(06)5 → clip(5-2)3。线性 epilogueBias融合保持精确数值等价。十二、归约优化decompose_reduction 与 rfactor12.1 decompose_reductiondecompose_reduction(block, loop) - SBlockRVschedule.py把一个归约 block 拆成两个init block由归约 block 的 init 语句翻译而来与update block去掉 init 后的原 block。init block 被插入到指定loop之前。前置条件输入是归约 blockloop是 block 的祖先loop不低于所有与归约 itervar 相关的循环。12.2 rfactorrfactor(loop, factor_axis) - SBlockRVschedule.py是归约并行化的经典变换把结合性归约按指定循环切分先生成n个分块的部分结果rf-block再累加回结果 bufferwb-block。factor_axis指定新维度插入位置范围[-ndim(B)-1, ndim(B)]负数按 numpy 约定归一化。rfactor 的完整前置条件源码 Note 部分包括loop只有一个子 block 且是归约 blockloop是归约循环loop未被 parallel/vectorize/unroll/bind所在作用域是 staged-pipeline归约 block 只写一个 bufferinit 与 body 都是简单BufferStore且模式注册为结合性归约预定义加、乘、min、max等。它把B[i] B[i] A[i, j, k]这类累加拆成可并行计算的部分和sch.rfactor(k, 0) # 生成 B_rf部分和 buffer与对应 rf-block、wb-block十三、Block 级标注storage_align / set_scope / unsafe_set_dtypestorage_align(block, buffer_index, axis, factor, offset)schedule.py要求第axis维 stride 满足stride k * factor offset例如factor2, offset1可避免 GPU shared memory 高维 bank 冲突。要求目标 buffer 是通过alloc_buffer定义的中间 buffer。示例中factor128, offset1最终让 buffer B 的 strides 变为[129, 1]set_scope(block, buffer_index, storage_scope)schedule.py把中间 buffer 的存储作用域改为shared/local等作用于alloc_buffer定义的 bufferunsafe_set_dtype(block, buffer_index, dtype)schedule.py修改中间 buffer 的数据类型如float32 → float16会插入T.cast。该原语不安全可能因类型转换改变程序正确性需谨慎使用。十四、blockize 与 tensorize向量化/硬件指令抽象14.1 blockizeblockize(target, preserve_unit_itersTrue) - SBlockRVschedule.py把以某循环为根的子树或指定的一组 block封装成一个新的外层 block要求给定循环下恰好有一个 block且 block 绑定能被该循环子空间整除。14.2 tensorizetensorize(block_or_loop, tensor_intrin, preserve_unit_itersTrue)schedule.py把循环包裹的计算替换为**张量指令tensor intrinsic**调用是接入硬件指令如tvm_mma_sync的关键T.prim_func(s_tirTrue) def mma_desc(a, b, c): ... # 描述型 intrinsic纯语义 T.prim_func(s_tirTrue) def mma_intrin(a, b, c): # 实现型 intrinsic真实指令 T.evaluate(T.tvm_mma_sync(...)) tvm.s_tir.TensorIntrin.register(test_mma_intrin, mma_desc, mma_intrin) sch tvm.s_tir.Schedule(before_tensorize) _, _, _, i1, _, _ sch.get_loops(sch.get_sblock(update)) sch.tensorize(i1, test_mma_intrin)tensorize 后内层 16×16×16 的循环体被替换为T.evaluate(T.tvm_mma_sync(...))同时生成A_1/B_1/C_1等按 tile 切分的match_buffer视图。14.3 tile_with_tensor_intrintransform.py 提供组合变换tile_with_tensor_intrin(sch, block, intrin_name, allow_paddingFalse) - LoopRV | None根据给定张量指令的 tile 形状自动对 block 的循环子集做 tilingallow_padding控制是否允许 padding找不到合法循环映射时返回None。intrin_name必须预先通过TensorIntrin.register注册。十五、通用标注annotate / unannotateannotate(block_or_loop, ann_key, ann_val)schedule.py给 block/loop 打键值对标注值支持str | int | float | ExprRV及其 list/dict 组合unannotate(block_or_loop, ann_key)schedule.py移除指定键的标注。标注会在 IR 中呈现为T.sblock_attr({...})供后续 pass 消费如 pipeline 优化提示。十六、布局变换transform_layout / transform_block_layout16.1 transform_layouttransform_layout(block, buffer, index_map, pad_valueNone, *, assume_injective_transformFalse)schedule.py用IndexMap变换 buffer 的布局sch.transform_layout( sch.get_sblock(B), buffer(write, 0), index_maplambda m, n: (m // 16, n // 16, m % 16, n % 16), )变换后 buffer 形状从(128, 128)变为(8, 8, 16, 16)所有访问自动改写为B[vi // 16, vj // 16, vi % 16, vj % 16]。pad_value用于指定变换引入 padding 区域的值可传int/float/Expr常量、IndexMap/Callable随变换索引变化的表达式如实现循环 bufferNone表示不允许引入 padding。源码特别提示若作用于输入 buffer调用方必须保证 pad 值实际存在否则后续代数化简等优化可能产生错误结果。assume_injective_transformTrue可跳过 index map 单射性检查对分析覆盖不到复杂 map 有用但单射性由调用者保证。16.2 transform_block_layouttransform_block_layout(block, index_map)schedule.py作用于 block 本身而非 buffer用IndexMap重排 block 的迭代空间如lambda i, j: (i * 16 j,)把二维 16×16 迭代空间折叠为一维 256。十七、Padding 分解与 Einsum 填充17.1 decompose_padding / can_decompose_paddingdecompose_padding(block, loop) - SBlockRVschedule.py把 padding 模式的计算 block 拆成两块常量 pad 填充 block填满整个写区域与in-bound 值填充 block仅 pad 谓词为真的区域前者插入到指定loop之前# beforey[vi] T.if_then_else(vi 6 and vi 134, x[vi - 6], 0) sch.decompose_padding(block, sch.get_loops(block)[0]) # afterblock_pad_const 填 0140 次block 仅做 y[vi 6] x[vi]128 次can_decompose_padding(block, loop) - boolschedule.py用于先检查 block 是否匹配 padding 模式。17.2 pad_einsumpad_einsum(block, padding)schedule.py把符合Einsum 模式buffer 访问下标只由常量或单个变量构成无复合表达式的 block 迭代域按因子向上取整补齐如 127→128并自动生成 producer/consumer padding block 避免越界。示例把A/B/C从(127, 127, 127)的 matmul pad 到(128, 128, 128)生成A_pad/B_pad/C_pad三个 buffer 及对应的拷贝 block。十八、滚动缓冲与杂项原语18.1 rolling_bufferrolling_buffer(block, write_buffer_index)schedule.py实现滚动缓冲沿祖先循环中带正向重叠的滚动轴折叠并循环复用 buffer同时给 block 追加谓词避免重复计算重叠元素。典型场景是 compute_at 级联后的滑动窗口算子如最大池化。前置条件包括block 非输出 block 且只有 RAW 依赖目标 buffer 是alloc_buffer定义的中间 bufferbuffer 生产者与消费者 LCA 是 for 循环访问区域至少一维含正向重叠。注意滚动后消费者 block 的region_cover属性会变为 false。18.2 其他杂项enter_postproc()schedule.py标记进入调度后处理阶段的无操作no-opunsafe_hide_buffer_access(block, buf_type, buf_index_array)schedule.py隐藏 block 的某些 buffer 访问buf_type为 read/writebuf_index_array为索引数组。不安全的原语可能破坏依赖分析典型用途是隐藏稀疏计算中的 indices buffer 访问以便后续 tensorizeannotate_buffer_access(block, buffer_index, buf_type, gen_new_ranges)schedule.py手工指定 block 的读/写区域。gen_new_ranges接收 block 的 iter_vars 并返回每个维度的Expr单点或(begin, end)元组。该功能在编译器无法准确推断复杂数据相关访问时非常有用会在 block 上添加explicit_read_region标注供CompactBufferAllocation等 pass 遵循标注错误会导致错误代码生成或运行时错误须确保区域覆盖所有真实访问。十九、底层机制Trace 与 sref 树调度之所以能做自动调优与决策回放依赖两套底层设施Tracetrace.py默认构造的TracedSchedule会把每次调度指令Instruction记录进轨迹配合各采样原语的decision参数即可在另一份调度上精确重放同一套决策。sch.copy()保证随机变量在拷贝中继续有效正是为派生多条候选调度而设计。SRef 树与依赖分析sblock_scope.pyStmtSRef是调度原语定位 IR 语句的统一句柄SBlockScope内以Dependency/DepKind记录 block 之间的读-写、写-读、写-写依赖支撑get_producers/get_consumers与大量原语的前置条件检查。二十、测试验证与进一步探索仓库中tests/python/s_tir/目录包含针对上述原语的完整测试集如test_tir_schedule_fuse.py、test_tir_schedule_rfactor.py等源码注释中fuse_reduction_epilogue也直接指向test_tir_schedule_fuse_reduction_epilogue作为示例参考。读者可结合以下路径继续深入核心实现python/tvm/s_tir/schedule/schedule.pySchedule全部原语、python/tvm/s_tir/schedule/state.pyScheduleState/ScheduleDebugMask、python/tvm/s_tir/schedule/transform.pytile_with_tensor_intrinC 侧实现与验证逻辑位于src/s_tir/schedule/目录Python 侧通过 python/tvm/s_tir/schedule/_ffi_api.py 绑定更上层的自动调度入口基于本模块的采样与 trace 机制位于 python/tvm/s_tir/meta_schedule/ 与 python/tvm/s_tir/dlight/。结语tvm.s_tir.schedule是 TVM 调度体系中最底层、最完整的命令式调度接口从Schedule构造参数seed/debug_mask/error_render_level/enable_check到数十个覆盖循环、缓存、计算位置、归约、布局、padding 的调度原语从随机采样与Trace回放到 sref 树与依赖分析本模块为手写高性能算子调度与自动调度器meta-schedule/dlight提供了统一而坚实的基础。掌握本指南中的原语语义与前置条件即可编写正确、可复现的 TensorIR 调度脚本并进一步理解 TVM 端到端编译流水线中的调度环节。赞分享模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载相关推荐Apache TVM S-TIR 变换 Pass 全解析从 tvm.s_tir.transform 到可调度 TensorIR 的编译管线Apache TVM S TIR 变换 Pass 全解析从 tvm.s_tir.transform 到可调度 TensorIR 的编译管线 S TIRSch模型编译深度学习推理引擎Apache TVM TensorIR 深度解析从张量程序抽象、TVMScript 编写到 DLight 与 MetaSchedule 自动化调度Apache TVM TensorIR 深度解析从张量程序抽象、TVMScript 编写到 DLight 与 MetaSchedule 自动化调度 Tenso模型编译深度学习推理引擎TVM TIRx Tile Primitives 完全指南可调度的 CUDA 硬件原语与分发机制TVM TIRx Tile Primitives 完全指南可调度的 CUDA 硬件原语与分发机制 本指南系统讲解 TVM 中 TIRx 方言的 Tile Pr模型编译深度学习推理引擎创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表