IBM Spyre 后端编译器开源:把 Inductor 的切分决定直接变成数据流流水线
核心概要
IBM 开源了面向 Spyre 加速器的数据流调度器,它读取 KTIR 这一 MLIR 中间表示,把 Inductor 已经定好的核间分工、片上暂存与访问切块转成数据流流水线,再经设备相关阶段生成 SpyreCode。此前这些决定缺少一个能承载完整数据并行视图的接口,后端需自行重新发现映射。首个算子接入比旧栈快约 4 倍,后续算子约 8 倍,目前覆盖带广播与归约维度的逐元素运算。
深度剖析
Spyre 的编译链现在有了一个明确的中间契约:前端 Inductor 决定哪个核算什么、数据放在哪块内存,后端只负责在这些决定之内安排数据流动。 此前 Inductor 的布局、切块、工作划分与暂存分配缺少一个能完整表达加速器数据并行视图的接口,后端难以直接承接这些决定。 KTIR 以核网格属性、显式内存空间和分布式内存视图逐项对应 Inductor 的每类决定,包括跨核切分的张量与核间重排。
调度器从一个最简单的合法流水线出发,逐步把它变成能真正在硬件上运行、并且更大的流水线。 起点是加载、计算、存储三段式流水线,但硬件并没有把设备内存经 FIFO 直连计算引擎的通路。 路径扩展按架构描述中的数据传输路径做最短路径搜索,把三段扩成五段,插入 DMA 载入、暂存载入、计算、暂存写回与 DMA 写回,并为每段标注执行单元。
流水线随后按存储层次逐层嵌套,并做双缓冲与核内实例切分,让搬运与计算重叠。 路径扩展后每个阶段仍按 64 元素向量执行一次,DMA 无法批量搬运,暂存也没有提前备好数据。 外层流水线按块粒度运行,内层按向量粒度运行;循环内跨阶段读写的暂存缓冲做双缓冲,核内实例数取自架构描述。
调度器的硬件知识全部放在一份架构描述文件里,换芯片只需换描述而不改调度器。 若把硬件细节硬编码进各个 pass,重定向到新设备就要改动调度器本身。 Spyre 有自己的描述,调度器测试跑在示例设备描述上,同一套 pass 在两者上不加修改地运行。
启示与展望
已演示的范围是带广播与归约维度的逐元素运算,覆盖多核 KTIR 程序,并从 KTIR 经调度与代码生成做了端到端验证,验证对象是代表性 transformer 运算。同一路径已从 Inductor 与 Triton 两个前端走通。矩阵乘法支持仍在完成中,成本模型驱动的性能优化也在加入,目标是端到端运行大语言模型。
V1 调度器中的循环结构、流水线粒度、缓冲深度、实例切分与块大小由局部启发式决定,而非搜索或成本模型;这些选择以占位符、取模因子和 distribute 计数的形式显式留在 IR 中,以便日后加入成本模型或自动调优而不改动流水线结构。算子接入提速倍数来自首个与后续算子的对比,未给出绝对性能数据。
