TileLang 软件流水线注解完全指南stage/order 手动调度与可重放 Bind 的边界【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang本文围绕 TileLang 的T.Pipelined注解体系展开从num_stages自动推断流水线到用stage/order数组手动调度 producer/consumer再到“可重放标量 Bind 为何不属于用户可见调度”这一核心设计决策。读完你可以独立完成非规则 GEMM 内核含额外后处理、重排序、手动异步拷贝分组的流水线注解编写并理解编译器在流水线重写中如何处理标量别名避免写出歧义或报错的注解。两种流水线写法num_stages 自动推断 vs stage/order 手动调度TileLang 可以从T.Pipelined(..., num_stages...)推断常见的 producer/consumer 流水线。对于常规 GEMM 类内核这是首选入口for ko in T.Pipelined(T.ceildiv(K, BK), num_stages3): T.copy(A[by * BM, ko * BK], A_shared) T.copy(B[ko * BK, bx * BN], B_shared) T.gemm(A_shared, B_shared, C_local)仓库中的 GEMM 示例 就采用了这种最简形式for k in T.Pipelined(T.ceildiv(K, block_K), num_stages3)。当循环体顺序不典型、存在额外后处理或需要手动异步拷贝分组时则改用显式注解for ko in T.Pipelined( T.ceildiv(K, BK), stage[0, 0, 1], order[0, 1, 2], ): T.copy(A[ko * BK], A_shared) T.copy(B[ko * BK], B_shared) T.gemm(A_shared, B_shared, C_local)从源码看两种模式最终汇聚到统一的深度推断逻辑。pipeline_utils.h 中GetPipelineNumStages()按三级优先级取流水线深度num_stages— 用户直接提供的级数tl_pipelined_num_stages— 由InjectSoftwarePipelinepass 回填s_tir::attr::software_pipeline_stage— 取max(stage) 1。这与文档给出的规则一致手动提供stage/order时不要再设num_stages深度由max(stage) 1推断num_stages单独用于编译器推断的流水线stage/order单独用于手动调度的流水线。T.Pipelined 接口 的 docstring 也明确编译器推断流水线中num_stages0表示不启用流水线手动调度时“prefer specifyingorderandstagewithoutnum_stages; the pipeline depth is inferred asmax(stage) 1”。接口层面T.Pipelined除stage/order外还暴露了sync同步元数据与groupproducer 分组元数据参数供手动流水线 lowering 使用def Pipelined( start, stopNone, num_stages0, orderNone, stageNone, syncNone, groupNone, annotationsNone, )用户模型只注解“做实际工作”的语句流水线注解描述的是可执行的流水线语句而非循环体内每个 IR 节点。应当注解的语句是那些产生实际效果effectful的操作拷贝包括 global-to-shared 与 TMA 拷贝填充fillGEMM 或其他 tile 操作归约reduction存储store与原子操作atomic当显式 wait、commit、同步语句属于手动流水线的一部分时。不应注解可重放的标量别名base: T.int32 ko * BK offset: T.int32 base tx这些Bind语句是值定义不是实体化的存储。编译器会在每个使用点自动放置它们。这一划分在编译器中有直接对应inject_pipeline.cc 中为每个被调度语句维护的PipelineAnnotation结构stage、order、async、async_group_id四个字段只作用于调度语句标量绑定走另一条 use-def 分析路径。stage 与 order两个数字与依赖检查规则每个被调度的语句携带两个数字stage逻辑流水线阶段数值越小越先执行order在同一阶段划分下发射语句的顺序。典型 copy/compute 流水线中拷贝在 stage 0、计算在 stage 1statement 0: copy A - stage 0, order 0 statement 1: copy B - stage 0, order 1 statement 2: gemm - stage 1, order 2stage与order两个数组按源码顺序与调度语句一一对齐。注解应用后编译器会做 buffer 依赖检查。若某调度语句为另一语句生产数据生产者不能被放在消费者之后。具体规则同一 stage 内生产者 order 必须小于消费者 order不同 stage 之间生产者 stage 必须小于等于消费者 stage。源码中对应的检查逻辑由ValidateScheduledBindDependencies与ValidatePipelineBody完成见 inject_pipeline.cc违反这些规则的注解会直接报错。重排序流水线order 先于源码顺序发射order数组的真正威力在于让流水线发射“未来迭代的生产者”先于“当前迭代的消费者”。例如for ko in T.Pipelined( num_tiles, stage[0, 1], order[1, 0], ): T.copy(A[ko * BK], A_shared) T.gemm(A_shared, B_shared, C_local)其含义是copy - stage 0, order 1 gemm - stage 1, order 0流水线重写时编译器会为每个阶段生成 prologue、稳态主体与 epilogue 片段并为每个 stage 配上正确的逻辑迭代号。源码顺序只决定哪个注解条目属于哪个调度语句order控制语句被归类之后的发射顺序。这一点可以从 测试文件 中大量stage[0,1]配order[1,0]的用例得到印证测试直接对T.serial循环挂注解并调用tl.transform.InjectSoftwarePipeline()验证重写结果。可重放标量 Bind为什么不能给它 stage/order这是本指南最核心的一条规则。可重放标量Bind可能出现在循环体中for ko in T.Pipelined( num_tiles, stage[0, 1], order[1, 0], ): base: T.int32 ko * BK T.copy(A[base tx], A_shared[tx]) T.gemm(A_shared, B_shared, C_local)注意stage和order只有两个条目对应 copy 与 gemmbase的定义被有意省略。这条规则之所以存在是因为可重放Bind没有唯一的流水线位置——它更接近 SSA 值或局部别名而非可执行操作。当同一标量被不同 stage 的语句使用时各消费者可能需要该值对应不同的逻辑流水线迭代for ko in T.Pipelined( num_tiles, stage[0, 1], order[1, 0], ): base: T.int32 ko * BK T.copy(A[base tx], A_shared[tx]) B[base tx] A_shared[tx]流水线重写后copy 可能需要“未来生产者迭代”的base而 store 需要“当前消费者迭代”的base。给base指定唯一 stage/order 只会导致越界或描述错误迭代。因此 TileLang 的处理模型是stage/order annotate only scheduled, effectful statements replayable Bind definitions are replayed at each consumer重放是 use-driven 的某消费者语句引用了流水线体中绑定的标量时编译器在该消费者之前重建所需的Bind并代入该消费者的逻辑访问索引。测试文件 中的test_inject_software_pipeline_replays_scalar_bind_dependencies专门验证了这一点重写后的 IR 中base/offset的绑定出现在 copy 与 store 两个消费者之前。读取 pipeline 外 buffer 的 Bind 仍可重放可重放标量Bind可以读取未在流水线体内被写入的 bufferidx: T.int32 Ids[ko] T.copy(A[idx], A_shared) B[idx] A_shared[0]此处idx仍是值别名。若两个调度语句都使用它TileLang 会按各自迭代重放idx Ids[logical_ko]。这可能重复一次Ids的加载但保留了Bind的别名语义值在消费者所处逻辑流水线迭代上计算。这与实体化生产者有本质区别idx_shared[tx] Ids[ko] T.copy(A[idx_shared[tx]], A_shared) B[idx_shared[tx]] A_shared[0]写入idx_shared创建了存储和真实的 producer/consumer 依赖。该语句应计入调度语句必要时像其他流水线 buffer 一样做版本化versioning。例外读取 pipeline 内 buffer 的 Bind 必须参与调度如果Bind读取的是同一流水线体内被写入的 buffer则不可重放T.copy(A[ko * BK], A_shared) value: T.float32 A_shared[tx] C[ko * BK tx] valueA_shared的加载依赖流水线生产者TileLang 会把这个Bind保留在调度语句列表中因此它需要一个stage与order条目。又因为标量Bind没有存储版本化能力被调度的Bind必须与使用其值的所有消费者处于同一 stage。如果你的本意是“加载一次、后续 stage 共享”应显式用 buffer/register 写入把值实体化而不是依赖Bind重放。Bind 依赖链先重放依赖者再重放使用点可重放 bind 可以依赖更早的可重放 bindbase: T.int32 ko * BK offset: T.int32 base tx T.copy(A[offset], A_shared[tx])编译器会在使用者之前按依赖顺序重放base offset consumer statement这保持了手动注解聚焦于真正的流水线操作同时保留了原始循环体中的词法标量依赖。旧式注解的向后兼容老代码可能在注解数组中包含可重放标量Bind的槽位for ko in T.Pipelined( num_tiles, stage[3, 0, 1], order[1, 2, 0], ): base: T.int32 ko * BK T.copy(A[base tx], A_shared[tx]) B[base tx] A_shared[tx]该形式出于兼容被接受标量Bind条目会被忽略其余条目按顺序落到调度语句上base Bind - ignored copy - stage 0, order 2 store - stage 1, order 0当旧式中某个可重放标量Bind被多个调度语句使用时TileLang 会发出警告。源码中可以直接找到这条警告逻辑inject_pipeline.ccif (multiple_consumers) { LOG(WARNING) Scalar Bind bind-var is used by multiple pipeline stages; its annotation is ignored and the bind is replayed at each use.; }新代码应优先使用省略可重放 bind 的更短注解数组仅在维护旧内核时才依赖旧式 bind 槽位。底层形式T.serial 直接挂注解大多数用户代码应使用T.Pipelined较低层的测试或 transform 级代码可以直接使用原始循环注解T.serial 接口 透传annotations字典for ko in T.serial( 0, num_tiles, annotations{ software_pipeline_stage: [0, 1], software_pipeline_order: [1, 0], software_pipeline_async_stages: [0], }, ): base: T.int32 ko * BK T.copy(A[base tx], A_shared[tx]) B[base tx] A_shared[tx]同样的规则适用software_pipeline_stage与software_pipeline_order只描述调度语句可重放标量Bind不需要条目。与异步拷贝相关的注解键定义见 pipeline_utils.h 中的software_pipeline_async_producers与software_pipeline_async_producer_groups测试文件中给出了完整组合示例software_pipeline_stage: [3, 0, 1]、software_pipeline_async_stages: [0]、software_pipeline_async_producers: [0, 1, 0]、software_pipeline_async_producer_groups: [-1, 0, -1]见 测试文件。设计动机稳定 API 与无歧义调度的权衡流水线 pass 把循环体拆成两个概念scheduled statements: executable operations controlled by stage/order replayable scalar bindings: local value aliases placed by use-def analysis这一拆分保证了用户侧 API 稳定并避免歧义调度。真正的流水线操作具有明确的执行点、可读写 buffer、可能需要异步拷贝记账或同步——用户为它指定 stage/order 是合理的。而可重放标量Bind不具备这些性质无副作用不拥有 buffer 存储可被多个调度语句使用其正确取值可能依赖消费者所处的逻辑流水线迭代。强制用户注解这类语句等于要求他们在“不存在唯一正确答案”时选择一个 stage/order。把标量 bind 在每个消费者处重放既在生成的 IR 中使别名语义显式化也让用户可见的调度聚焦于真实的流水线工作。编写手动流水线注解的自检清单来自 原始指南 的完整 checklist逐条核对构建stage与order时只数调度语句省略可重放标量别名如base ko * BK、idx Ids[ko]前提Ids未在流水线体内被写入若Bind读取了流水线体内被写入的 buffer将其计入调度语句保持每个调度语句的order唯一按 stage/order 依赖规则保持生产者先于消费者同 stage 比 order跨 stage 比 stage;software_pipeline_async_producers与software_pipeline_async_producer_groups的长度与调度语句列表一致新代码优先使用无 bind 的注解形式旧式 bind 槽位仅用于维护旧内核。延伸阅读路径文档主体software_pipeline.mdPython 接口T.Pipelined / T.serial核心变换 passInjectSoftwarePipeline、深度推断 GetPipelineNumStages变换级测试含旧式注解、Bind 重放、依赖重放用例test_tilelang_transform_Inject_software_pipeline.py自动推断流水线的另一入口pipeline_planning.cc适用前提本文所有行为以当前仓库的src/transform/inject_pipeline.cc实现为准注解键名software_pipeline_stage等属于 transform 级接口随编译器内部演进可能变化用户层代码应优先使用T.Pipelined的高层参数而非直接挂注解字典。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
