From KTIR to SpyreCode: An MLIR Dataflow Backend for IBM Spyre

从 KTIR 到 SpyreCode：IBM Spyre 的 MLIR 数据流后端 - PyTorch

特色项目

TL;DR

Torch-spyre 的 Inductor 决定每个内核的工作如何在 Spyre 的核心之间分配，并管理张量在片上缓存和片外内存之间的驻留状态，而后端编译器将这些决策转化为硬件实际可执行的内容。它们之间的契约通过 KTIR（内核分块中间表示）来表示，这是一种基于加速器数据并行视图的 MLIR 表示，其中每个计算分块都能访问所有内存，单个逻辑张量可以分布在多个核心的缓存中。新的开源数据流调度器消费 KTIR，并将每个分块的工作转化为数据流管道：它将内核拆分为程序，从简单的加载 → 计算 → 存储管道开始，根据架构描述中的内存层次路由所有传输，并将该管道扩展为嵌套的、双缓冲的、多实例的调度。调度器的输出使用 DataflowIR 表示，这是一种 MLIR 中间表示，Spyre 的设备特定阶段使用它生成 SpyreCode，即一个符号化的 SpyreBinary 加上一个作业计划，该计划由 Spyre 运行时一次性准备并多次启动，在启动时填充张量地址和形状。

本文是三篇系列文章的中篇。《为 IBM Spyre 加速器扩展 Inductor》介绍了 torch-spyre 的 Inductor 扩展如何确定布局、分块、工作划分和缓存放置，《将 Spyre 构建为原生 PyTorch 设备》描述了运行时如何启动一个编译后的工件，其符号值仅在最后一刻填充。本文将重点介绍 Inductor 做出的决策如何实际转化为运行时可启动的程序。

后端编译器在 Spyre 软件栈中的位置

设计中贯穿三个核心理念：

• 前端决定并行分解，后端负责执行。KTIR 的设计用于承载 Inductor 已经选择的映射（哪个核心计算什么数据，从哪个内存获取），这样后端无需重新发现这些信息。后端的任务是在该映射内进行数据流调度。

• 调度意味着构建流水线。Spyre 是一个具有解耦访问和执行的流水线机器，因此加载、存储和计算在独立引擎上通过 FIFO 进行通信。调度器从最简单的合法流水线开始，逐步进行一次转换优化，每一步都是可读的 MLIR 重写。

• 边界是 IR 而非库函数调用。KTIR、调度器的 ScheduleIR 和 DataflowIR 都是基于上游方言（linalg、scf、arith、memref、tensor）的普通 MLIR，调度器及其方言均以 Apache-2.0 开源。设备特定的代码生成阶段从 DataflowIR 边界开始。

KTIR：Inductor 与后端之间的契约

为何需要新接口

高效后端调度的关键要求之一是 torch-spyre 的 Inductor 与后端编译器之间具备完整数据并行视图的契约，该视图能够捕捉所有计算分块都能访问所有内存、单个逻辑张量可以分布在多个核心缓存中的加速器特性。

KTIR（RFC 0682，实现于torch-spyre/ktir-mlir-frontend）是一种基于开源MLIR的基于瓦片（tile）的、块结构的中间表示（IR）。我们的目标是使其表达能力足够强大，能够捕捉来自两个截然不同来源的数据并行映射：一方面来自Inductor，另一方面来自用Triton或Helion等语言编写的Spyre内核。

对加速器的数据并行抽象

KTIR基于对硬件的刻意简化模型进行设计：一组核心，每个核心包含一个计算引擎和自己的片上缓存，通过一个片上互连网络连接，该网络还可连接一个或多个片外内存。最关键的特点是每个计算单元都能全局访问所有内存，包括自己的缓存、其他核心的缓存以及设备内存。

KTIR中多核加速器的硬件抽象（来自RFC 0682）

该模型与Spyre的架构高度契合：环形排列的32个核心，每个核心配备2MB片上缓存，由片外设备内存支持。数据通过固定大小的128字节数据块（称为"sticks"）在设备内存和缓存之间传输，每个数据块存储64个fp16值。刻意省略了FIFO、加载/存储单元和解耦执行引擎等细节。KTIR描述每个计算单元执行的操作及其数据所在位置，但不涉及数据如何到达那里，因为确定数据传输路径正是调度器的职责。

KTIR还具有组合性。算术运算、归约和控制流直接来自上游MLIR方言（arith、math、linalg、scf、tensor、memref），而专用的ktdp方言则补充了其他方言无法表达的两个关键要素：显式内存和显式访问。另一个规模较小的spyreop方言则覆盖了Spyre特有的内置函数。ktdp的核心操作包括：

ktdp

操作

get_compute_tile_id

标识当前计算单元；可以是1维或多维，基于内核的

grid

属性

construct_memory_view

将某个（可能是符号化）地址处的命名内存空间中的区域解释为一个步长张量；不涉及分配，也不涉及数据移动

construct_distributed_memory_view

将多个覆盖全局坐标空间不同部分的内存视图组合成一个逻辑张量

construct_access_tile

对视图的结构化坐标集：用整数集定义域，用仿射映射定义遍历顺序，使瓦片可以是矩形、步长、倾斜或三角形

construct_indirect_access_tile

类似，但某些维度通过另一个张量进行索引，例如分页KV缓存收集

load

/

store

唯一的数据移动方式：消耗一个访问瓦片并生成或写入数据瓦片

inter_tile_produce

inter_tile_reduce

计算单元之间的通信，例如部分和的全归约

以下是一个完整的内核示例：对分布在32个计算单元（每行3个单元）中的96×64 fp16张量执行逐元素加法。

#tile = affine_set<(d0, d1) : (d0 >= 0, -d0 + 2 >= 0, d1 >= 0, -d1 + 63 >= 0)>   // 3x64
#full = affine_set<(d0, d1) : (d0 >= 0, -d0 + 95 >= 0, d1 >= 0, -d1 + 63 >= 0)>
#id   = affine_map<(d0, d1) -> (d0, d1)>

// 张量的起始地址作为参数传入：PyTorch 的分配器拥有这些地址。
func.func @add(%A: index, %B: index, %C: index) attributes {grid = [32]} {
  %c0   = arith.constant 0 : index
  %c3   = arith.constant 3 : index
  %id   = ktdp.get_compute_tile_id : index
  %row0 = arith.muli %id, %c3 : index                 // 此 tile 的 3 行块

  %A_view = ktdp.construct_memory_view %A, sizes: [96, 64], strides: [64, 1] {
      coordinate_set = #full, memory_space = #ktdp.memory_space<global>
  } : memref<96x64xf16>
  // ... %B_view 和 %C_view 类似

  %A_t = ktdp.construct_access_tile %A_view[%row0, %c0] {
      access_tile_set = #tile, access_tile_order = #id
  } : memref<96x64xf16> -> !ktdp.access_tile<3x64xindex>
  // ... %B_t 和 %C_t 类似

  %a = ktdp.load %A_t : !ktdp.access_tile<3x64xindex> -> tensor<3x64xf16>
  %b = ktdp.load %B_t : !ktdp.access_tile<3x64xindex> -> tensor<3x64xf16>
  %e = tensor.empty() : tensor<3x64xf16>
  %c = linalg.add ins(%a, %b : tensor<3x64xf16>, tensor<3x64xf16>)
                  outs(%e : tensor<3x64xf16>) -> tensor<3x64xf16>
  ktdp.store %c, %C_t : tensor<3x64xf16>, !ktdp.access_tile<3x64xindex>
  return
}

该内核将四个关注点分离。视图定义了内存的解释方式，访问块指定了某个 tile 所涉及的坐标，加载和存储是唯一移动数据的操作，而 linalg 负责算术运算。请注意工作划分已经确定：grid = [32] 与 %id 的算术运算共同构成了划分方式，后端应遵循这一划分而非自行猜测。

Spyre 内置函数：spyreop 方言

上游的 arith 和 math 覆盖了内核所需的大部分算术运算，但 Spyre 的计算单元还提供了某些操作，这些操作在任何通用方言中都没有对应的名称。如果内核需要这些操作，应能直接请求而非依赖后端识别特定的算术操作模式。为此，KTIR 提供了第二个更小的方言 spyreop，将 Spyre 内置函数作为普通 MLIR 操作暴露出来。它涵盖了逐元素运算和激活函数（exp、gelu、sigmoid、silu、softplus、sqrt、rsqrt、reciprocal、realdiv）、归一化构建块（联合均值和均方值、LayerNorm 的归一化和缩放部分，部分支持融合变体）、整数和地址算术（包括 idx32toaddr，用于计算 base + stride × index 的 gather 操作），以及比较和选择操作。该方言还定义了 df16，这是 IBM 的 16 位 DLFloat 格式，其符号位、指数位和尾数位分配为 1-6-9，而非 IEEE 半精度的 1-5-10，同时还定义了“融合”类型，将两个浮点数作为单个值进行处理。

内置函数出现在任何其他标量计算可能的位置，即在 linalg.generic 的主体内部：

%y = linalg.generic {indexing_maps = [#in, #out],
                     iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
       ins(%x : tensor<6x64x64xf16>) outs(%init : tensor<1x6x64x64xf16>) {
  ^bb0(%a: f16, %out: f16):
    %g = spyreop.gelu %a : f16
    linalg.yield %g : f16
} -> tensor<1x6x64x64xf16>

因为它们只是线性代数体中的操作，内在函数会像其他任何操作一样与周围的逐元素工作融合，调度器会将它们作为计算步骤传递，直到准备将它们降低。每个内在函数如何扩展为硬件自身的指令序列，完全不会在调度器中编写：它存在于架构描述中，作为本文后面描述的设备模式之一。

KTIR 如何捕捉 Inductor 的工作映射

Inductor 的文章介绍了在生成任何代码之前，torch-spyre 做出的一系列决策。每个决策在 KTIR 中都有直接对应项：

torch-spyre 的 Inductor 传递中的决策

在 KTIR 中的表现形式

工作划分：每个操作的核拆分，例如矩阵乘法的 16 × 2 网格

内核的

属性，加上对

ktdp.get_compute_tile_id

的仿射算术，选择每个分块的切片

棒状粒度布局：哪一维是棒状，其他维度的分块顺序

的步长，以及每个访问分块的

access_tile_order

粗粒度分块和工作集缩减

围绕加载、计算和存储的

scf.for

循环，逐个分块执行链式操作

由 scratchpad 分配器选择的 scratchpad 居住位置

在分配器分配的偏移处的

#ktdp.memory_space<ct_local, ct_id = k>

内存视图

跨核拆分张量，使每个切片适合 scratchpad

每个核的

ct_local

视图上的张量拆分

核到核的重新布局（

lx_relayout

）

对一个分布式视图的访问分块，存储到具有不同分布的视图中

跨核的缩减（一个输出分块的局部和）

分页张量和收集操作

设备特定的激活和归一化步骤

spyreop

内在函数嵌入在

linalg.generic

体中

由 PyTorch 分配器拥有的张量地址

内核参数。它们一直保持符号化直到 SpyreCode，并在启动时由运行时替换

仅在运行时已知的张量形状，在一个边界内

!ktdp.runtime_arg<index, granularity = G, upperbound = U>

，保持符号化并在启动时以相同方式替换

最后两行是本文与运行时文章的交汇点。因为 KTIR 允许起始地址或大小是值而非常量，所有下层都可以保持该值符号化，直到运行时能够提供它。

针对核分布的 scratchpad 中固定的张量

KTIR 与基于 GPU 的 tile IR 最大的不同之处在于其对片上内存的处理方式。像 Triton 这样的 tile 语言会描述每个程序实例从全局内存加载张量块，对其进行计算，然后将结果存储回去。这些块如何通过共享内存、寄存器和缓存进行调度，由编译器和硬件决定，而程序实例的片上存储空间对其私有且仅在其运行期间存在。在这种模型中，通常没有方法，也几乎没有必要说明当下一个操作读取张量时，该张量应该仍然保留在某个核心的片上内存中，或者一个程序实例应该读取另一个程序实例在片上存储的数据。这种模式在具有缓存和动态调度器的硬件上运行良好，但不适用于 Spyre。Spyre 没有硬件管理的缓存可供回退，而将张量从一个操作保留到下一个操作是最重要的性能杠杆之一。在 Spyre 上，Scratchpad 不是内核偶然命中的缓存，而是编译器选择使用的内存，Inductor 的 Scratchpad 分配器决定哪些缓冲区驻留、驻留地址以及驻留的核心。因此，KTIR 将片上内存的放置作为 IR 的核心组成部分，而不是实现细节：它直接记录这一决策。

每个内存空间对应物理内存：#ktdp.memory_space<global> 是设备内存，#ktdp.memory_space<ct_local, ct_id = k> 是计算单元 k 的 Scratchpad。视图的 coordinate_set 记录它持有更大逻辑张量的哪一部分，因此如果你将每个核心的视图与 construct_distributed_memory_view 结合使用，你会得到一个以全局坐标索引的单一张量，其存储分布在多个 Scratchpad 中。

// X[128, 1024] f16，由内核的早期操作留在 Scratchpad 中：
// 每个核心的 32 行，位于 Scratchpad 分配器选择的偏移位置。
#rows0 = affine_set<(d0, d1) : (d0 >= 0,      -d0 + 31 >= 0,  d1 >= 0, -d1 + 1023 >= 0)>
#rows1 = affine_set<(d0, d1) : (d0 - 32 >= 0, -d0 + 63 >= 0,  d1 >= 0, -d1 + 1023 >= 0)>
// ... #rows2, #rows3

%X0 = ktdp.construct_memory_view %lx_off, sizes: [32, 1024], strides: [1024, 1] {
        coordinate_set = #rows0, memory_space = #ktdp.memory_space<ct_local, ct_id = 0>
      } : memref<32x1024xf16>
%X1 = ktdp.construct_memory_view %lx_off, sizes: [32, 1024], strides: [1024, 1] {
        coordinate_set = #rows1, memory_space = #ktdp.memory_space<ct_local, ct_id = 1>
      } : memref<32x1024xf16>
// ... %X2, %X3
%X = ktdp.construct_distributed_memory_view (%X0, %X1, %X2, %X3
        : memref<32x1024xf16>, memref<32x1024xf16>, memref<32x1024xf16>, memref<32x1024xf16>)
        : memref<128x1024xf16>

// 计算单元 1 读取第 48-79 行：一半来自自己的 Scratchpad，一半来自核心 2 的。
#rows32 = affine_set<(d0, d1) : (d0 >= 0, -d0 + 31 >= 0, d1 >= 0, -d1 + 1023 >= 0)>
%id  = ktdp.get_compute_tile_id : index
%row = affine.apply affine_map<(d0) -> (d0 * 32 + 16)>(%id)
%t = ktdp.construct_access_tile %X[%row, %c0] {
       access_tile_set = #rows32, access_tile_order = #id
     } : memref<128x1024xf16> -> !ktdp.access_tile<32x1024xindex>
%x = ktdp.load %t : !ktdp.access_tile<32x1024xindex> -> tensor<32x1024xf16>

分布式内存视图：一个逻辑张量，跨四个 Scratchpad 固定

这一构造完成了大量工作。首先，它允许中间张量在连续操作之间保留在芯片上：生产者可以将数据写入分布式ct_local视图，消费者可以从该视图中读取数据，而无需涉及设备内存，这正是Inductor文章中softmax示例所依赖的那种无需往返的融合方式。其次，它将核心划分不匹配转换为普通IR。当生产者和消费者以不同方式划分张量时，消费者的访问分块只需跨越缓存边界，由于每个计算分块都能访问所有缓存，这仍然只是一个加载操作；torch-spyre协同优化分配器在权衡溢出时需要的片上重排操作无需专用指令。第三，它保持了放置绑定。运行时后端强调"放置必须显式且绑定，因为编译代码可能依赖于此"，而KTIR正是记录这种依赖性的位置，因为视图上的内存空间和坐标集构成了放置契约。

后端编译器概览

后端编译器是一个包含两个阶段的单一流水线。开源的数据流调度器通过ScheduleIR将KTIR转换为DataflowIR，随后三个设备专用阶段（符号化SpyreBinary生成、运行时符号替换元数据生成、SpyreCode生成）将DataflowIR转换为SpyreCode。DataflowIR是两个阶段交汇的唯一节点。

Spyre后端编译器：KTIR输入，SpyreCode输出

文章其余部分将按照图示从左到右展开，重点讲解调度器，然后简要覆盖其下方的三个阶段。

数据流调度器：KTIR → ScheduleIR → DataflowIR

数据流加速器是由宽数据路径（含FIFO）连接的分布式计算引擎和缓存内存集合。它采用解耦的访问-执行设计，其中数据移动（DMA、缓存读写）和计算操作在独立引擎上并行执行。在这样的机器上，调度意味着构建跨越空间和时间的数据流管道，决定每个步骤由哪个引擎执行以及步骤之间如何重叠。

数据流调度器通过MLIR传递流水线实现这一目标，其构建遵循几个核心原则。它基于架构描述而非硬编码到特定芯片，下面将更详细地探讨这一点。每个调度都基于三个构造构建的单一执行模型：嵌套的scf.for循环、ktdf.pipeline（跨迭代重叠的阶段）和ktdf.parallel（在迭代空间中分割的相同实例），所有分析和转换都基于这些术语进行。只要可能，调度器会依赖上游MLIR，将计算转换为linalg.generic并将其与上游传递融合，使用scf循环和memref缓冲区。整个过程遵循一个策略：从最原始的正确流水线开始，然后让后续传递逐步在硬件上使其合法化或逐步扩展，每一步都可以独立测试。

传递分为三组，每组在边界处都有固定的IR契约，因此可以在组内添加或重新排序传递而不会影响其他组。

数据流调度器：三个传递组，每组一个IR契约

在四个命名的检查点（pre_mapping、pre_scheduling、post_scheduling 和 post_lowering）中，调度器会应用用PDL编写的重写模式。这些模式存储在架构描述中而非调度器本身，这意味着设备可以提供自己的lowerings而无需fork调度器。这就是每个spyreop内在函数（如spyreop.exp或spyreop.gelu）被替换为硬件指令模板的地方，也是内维度归约变为SIMD归约的位置。

调度器的中间形式ScheduleIR是与linalg、scf和memref一起使用的ktdf方言。其词汇量足够小，可以完整列出：

ktdf

构造

含义

ktdf.pipeline [modulo(size: N)]

跨迭代重叠的阶段链；

modulo

在N个资源副本之间轮转（N缓冲）

ktdf.stage depends_in(...) depends_out(...)

流水线的一个阶段，最终绑定到某种执行单元（

applicable_units

ktdf.private

专属于一个流水线的资源：FIFO槽、Scratchpad缓冲区、令牌

!ktdf.token

阶段间纯粹的顺序值，自身没有运行时开销

!ktdf.fifo.slot<SRC -> DST, 64xf16>

两个单元之间硬件FIFO中的一个槽位

ktdf.data_transfer

ktdf.ind_data_transfer

显式在内存和FIFO之间移动数据，直接或通过间接寻址缓冲区

ktdf.read_from_fifo

ktdf.write_to_fifo

计算阶段如何消费和生成FIFO数据

ktdf.parallel ... distribute(num_instances = N)

将迭代空间分割为N个相同硬件实例上的连续块

ktdf.tiling.reserve_size

derive_size

linearize_index

后续选择的分块大小占位符、当前分块的范围和原始索引

架构描述

调度器对硬件的所有了解都来自一个用ktdf_arch方言编写的MLIR文件。该文件将设备描述为一个图。其节点是内存（包含类型和大小），以及执行单元（包含调度器关注的特性：每元素类型的SIMD通道数、加载/存储单元的字长和访问粒度、每个FIFO的深度）。其边是数据路径，每条边都有传输粒度。组通过类似硬件的方式对资源进行作用域划分，例如一个拥有Scratchpad的核，以及使用该Scratchpad的加载/存储和计算单元。以下是调度器测试中使用的示例设备的简化版本：

#L1  = { kind = "L1", size = 1048576 }                           // 1 MB scratchpad
#SFU = { kind = "SFU", ktdf_arch.features = {
           ktdf_arch.feature.compute,
           ktdf_arch.feature.simd = { lanes = #ktdf_arch.map<f16 = 64, f32 = 32, i32 = 32> } } }
#L1LU_CORE_FIFO = { ktdf_arch.features = { ktdf_arch.feature.queue = { depth = 16, ordered } } }

ktdf_arch.device @sample_device {
  %ddr = memory #DDR
  group { kind = "core" } share(%ddr) {
    %l1 = memory #L1
    %mnilu = group { kind = "MNILU_Block" } share() { ... } -> exec_unit   // DMA load unit
    %mnisu = group { kind = "MNISU_Block" } share() { ... } -> exec_unit   // DMA store unit
    datapath #DDR_MNILU %ddr to %mnilu : memory, exec_unit
    datapath #MNILU_L1 %mnilu to %l1 : exec_unit, memory
    %sfu  = group { kind = "SFU_Block" } share() { ... } -> exec_unit   // compute unit
    %l1lu = exec_unit #L1LU                                              // scratchpad load unit
    %l1su = exec_unit #L1SU                                              // scratchpad store unit
    datapath #L1LU_CORE_FIFO %l1 to %l1lu : memory, exec_unit
    datapath %l1lu to %sfu : exec_unit, exec_unit
    ...
  }
  patterns ["pre_scheduling"] { ... }                   // device-specific rewrites, in PDL
  patterns ["post_scheduling"] { ... }
}

内核通过一行声明指定目标，该声明导入此文件并将KTIR的内存空间映射到设备的内存上，例如将global映射到DDR，ct_local映射到L1。从那时起，流水线中的几乎所有决策都会从描述中读取信息。SIMD通道数量决定了第一条流水线的宽度。数据路径构成了路径扩展搜索的路由图，因此传输所经过的跳数取决于图的允许范围。访问粒度、传输粒度和FIFO深度限制了每次传输的形状。内存大小限制了地址分配的可能范围，从而决定了哪些分块尺寸可以适配。而模式块则提供了设备特有的降级方式，例如在调度器的检查点处，每个spyreop内在函数如何扩展为硬件指令。

由于所有这些信息都保存在描述中而非传递过程中，重新定位调度器意味着编写新的描述，而非修改调度器本身。Spyre拥有自己的描述，调度器的测试针对上述示例设备运行，相同的传递过程在两种设备上均无需修改即可运行。

第一步：将内核拆分为多个程序

KTIR内核是一系列操作的序列，而Spyre程序是加载到设备上的单一流水线配置，因此调度器的第一个实际决策是将内核切割为可单独调度的多个程序。

调度器通过从每个ktdp.store反向追踪，经过为其提供数据的计算部分，最终追溯到其依赖的ktdp.load，以此确定切割位置。所有通过这种方式连接的部分形成一个计算组，即一条通过内存的单个加载→计算→存储链。一个组的结束位置要么是属于不同组的加载操作，要么是任何包含区域的操作（如循环或条件语句）。每个组随后会成为自己的嵌套MLIR模块，即module @local_schedule_N，而原始内核则会转换为一个协调器，保留控制流和地址运算，并按顺序调用这些程序。因此程序数量正好等于计算组的数量。

// KTIR：与循环嵌套交错进行加载、计算和存储
func.func @kernel() {
  scf.for %i = %c1 to %c5 step %c1 {
    ... ktdp.load A ... linalg.add ... ktdp.store B ...
    scf.for %j = %c1 to %c3 step %c1 {
      ... ktdp.load B ... linalg.mul ... ktdp.store C ...
    }
    ... ktdp.load C ... linalg.sub ... ktdp.store D ...
  }
}

// ScheduleIR：一个协调器，它为每个计算组调用一个程序
module {
  func.func @kernel() {
    scf.for %i = %c1 to %c5 step %c1 {
      func.call @local_schedule_1() : () -> ()
      scf.for %j = %c1 to %c3 step %c1 {
        func.call @local_schedule_0() : () -> ()
      }
      func.call @local_schedule_2() : () -> ()
    }
    return
  }
  func.func private @local_schedule_0()
  ...
}
module @local_schedule_0 { func.func @local_schedule_0() { /* views, tile id, load, mul, store */ } }
module @local_schedule_1 { ... }
module @local_schedule_2 { ... }

这种拆分带来的一个后果值得特别说明，因为它将两种截然不同的循环类型分离开来。出现在KTIR中的循环（通常来自Inductor的粗粒度分块）保留在协调器中，成为按分块顺序启动程序的序列。而程序内部的循环则属于调度器：它通过将融合计算分块到硬件的向量宽度来创建这些循环，并可以自由地以任何方式重构它们。

每个程序都会从内核中获取自己需要的副本，包括内存视图、常量、ktdp.get_compute_tile_id和网格属性，因此Inductor的工作划分可以完整地传递到每个程序。具有间接访问的内核会进一步拆分：第一个程序在间接地址缓冲区中将索引转换为地址（使用spyreop.idx32toaddr），第二个程序则使用该缓冲区执行gather或scatter操作。最后，在每个程序内部，逐元素操作会被转换为linalg，泛化为linalg.generic，并沿着生产者-消费者链进行融合，使每个程序最终只保留一个融合后的linalg.generic，这是调度器其余部分工作的基本单元。

第2步：最简单的可行流水线

接下来，调度器为每个程序构建其第一个数据流流水线。它将融合后的通用计算分块到默认计算单元的SIMD宽度（从架构描述中读取，示例设备为64条fp16通道），将每个生成的循环标记为parallel_loop或reduction_loop，并在最内层循环中构建一个三阶段的ktdf.pipeline，该流水线从内存加载到FIFO，执行计算，再从FIFO存储回内存。

scf.for %i ... { scf.for %j ... {
  ktdf.pipeline {
    %p:6 = ktdf.private -> (!ktdf.fifo.slot<"DDR" -> "SFU", 64xf16>, !ktdf.fifo.slot<"DDR" -> "SFU", 64xf16>,
                            !ktdf.fifo.slot<"SFU" -> "DDR", 64xf16>, !ktdf.token, !ktdf.token, !ktdf.token) {...}
    ktdf.stage depends_in(none) depends_out(%p#3) {
      ktdf.data_transfer from %A[%i, %j] size [1, 64] to %p#0 size [1, 64] : ...
      ktdf.data_transfer from %B[%i, %j] size [1, 64] to %p#1 size [1, 64] : ...
    }
    ktdf.stage depends_in(%p#3) depends_out(%p#4) {
      %a = ktdf.read_from_fifo %p#0 : <"DDR" -> "SFU", 64xf16> -> tensor<1x64xf16>
      %b = ktdf.read_from_fifo %p#1 : <"DDR" -> "SFU", 64xf16> -> tensor<1x64xf16>
      %r = linalg.generic {...} ins(%a, %b : ...) outs(%e : tensor<1x64xf16>) {...}
      ktdf.write_to_fifo %r, %p#2 : tensor<1x64xf16>, <"SFU" -> "DDR", 64xf16>
    } {applicable_units = ["SFU"]}
    ktdf.stage depends_in(%p#4) depends_out(%p#5) {
      ktdf.data_transfer from %p#2 size [1, 64] to %C[%i, %j] size [1, 64] : ...
    }
  }
} {loop_type = #ktdf.loop_type<parallel_loop>} } {loop_type = #ktdf.loop_type<parallel_loop>}

（本节中的代码片段摘自调度器的lit测试，这些测试使用了一个示例架构，其内存称为DDR和L1，计算单元称为SFU。在Spyre上，设备内存、每个核心的Scratchpad以及计算流水线分别承担相同的角色。）

该流水线在功能上是完整的，但目前仍不合法，因为硬件中没有任何部分能通过FIFO将设备内存直接连接到计算引擎。修复这个问题是下一组处理步骤的首要任务。

第三步：路径扩展通过内存层次路由数据

路径扩展通过架构描述使数据移动合法化。它从设备构建一个路由图，将内存、执行单元和加载/存储单元作为节点，数据路径作为边，然后针对每个流水线，从三阶段版本（如设备内存→计算→设备内存）推导出端点，对每对连续端点运行最短路径搜索，并将这些跳转连接起来。当实际路径比端点建议的路径更长时，该处理步骤会重写流水线，为每个中间跳转插入传输阶段，在Scratchpad中分配中转缓冲区，重新定义FIFO以确保每个FIFO连接两个实际相邻的单元，并为每个阶段标记将要执行的单元类型。

在示例中，三个阶段扩展为五个阶段：DMA加载单元将一个tile加载到Scratchpad，Scratchpad加载单元将其送入FIFO，计算引擎进行处理，Scratchpad存储单元从结果FIFO中取出数据，DMA存储单元将结果写回设备内存。

ktdf.pipeline {
  %p:10 = ktdf.private -> (memref<1x64xf16, "L1">, memref<1x64xf16, "L1">,
                           !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16>, ..., !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>,
                           memref<1x64xf16, "L1">, !ktdf.token, ...) {...}
  ktdf.stage ... { ktdf.data_transfer from %A[%i, %j] ... to %p#0[0, 0] ... }  {applicable_units = ["MNILU"]} // DDR -> L1
  ktdf.stage ... { ktdf.data_transfer from %p#0[0, 0] ... to %p#2 ... }         {applicable_units = ["L1LU"]}  // L1 -> FIFO
  ktdf.stage ... { ... read_from_fifo ... linalg.generic ... write_to_fifo ... } {applicable_units = ["SFU"]}   // compute
  ktdf.stage ... { ktdf.data_transfer from %p#4 ... to %p#5[0, 0] ... }         {applicable_units = ["L1SU"]}  // FIFO -> L1
  ktdf.stage ... { ktdf.data_transfer from %p#5[0, 0] ... to %C[%i, %j] ... }  {applicable_units = ["MNISU"]} // L1 -> DDR
}

这也是KTIR中的布局信息开始发挥作用的地方。如果某个操作数已经是ct_local视图，意味着Inductor已将其固定在临时存储区，那么它到计算引擎的路径将直接从临时存储区开始，路径扩展过程完全不会为其添加任何DMA阶段。管道直接从临时存储区到FIFO再到计算单元，张量无需经过设备内存的往返。换句话说，Inductor中决定的驻留策略会直接转化为调度器中的管道结构，中间无需任何额外机制。

一些后续的合法性检查补充完善了这一过程：间接地址缓冲区填充通过填充单元实际可读的内存进行分阶段处理，标量操作数会被扩展为完整的向量宽度，并作为硬件 splat 进行传输。

第4步：逐层构建嵌套管道，按内存层次结构逐级处理

路径扩展后，每个阶段仍按每个64元素向量运行一次。这在语法上是合法的，但效率低下：DMA引擎通过批量传输最高效地移动数据，而临时存储区的核心作用就是让数据在计算需要之前就已就绪。因此调度器会将程序的循环嵌套重塑为管道的层次结构，每个管道按其所在内存层次结构的粒度运行。

它首先将并行循环分块为分块循环和点循环，但将每个分块大小作为占位符保留在IR中（ktdf.tiling.reserve_size），并附上约束条件，以便后续选择具体大小。然后确定哪些阶段必须同步运行，因为它们通过FIFO或寄存器通信，而哪些阶段不需要同步，例如仅访问设备内存和临时存储区的DMA阶段。对分块的循环会下沉到这些组中。DMA阶段会将整个分块一次性传输到临时存储区，而临时存储区加载→计算→临时存储区存储的链则会合并为一个阶段，其中包含其自身的嵌套ktdf.pipeline。最终得到的结构是一个外层管道，按分块逐个处理，包裹着一个内层管道，按向量逐个处理。

非所有新管道都是嵌套的。在循环迭代之间不会发生变化的传输，例如沿该循环维度广播的操作数或常量缓冲区，会被提升到一个同级的管道中，在循环之前运行一次，而不是在每次迭代时重新获取相同的数据。归约操作则以类似精神进行重构：较大的归约会被分割为适合缓存的块，外层部分会变成带有累积器的显式循环，而最内层的、与硬件宽度匹配的部分则留给设备模式映射到硬件自身的归约支持上。

scf.for %m0 ... { scf.for %n0 ... {
  ktdf.pipeline {                                            // outer: tile granularity
    %p0:5 = ktdf.private -> (memref<2x1x64xf16, "L1">, ..., !ktdf.token, !ktdf.token) {...}
    ktdf.stage depends_in(none) depends_out(%p0#3) {         // bulk DDR -> L1 for the tile
      scf.for %m1 ... { scf.for %n1 ... { ktdf.data_transfer from %A[..] ... to %p0#0[..] ... } }
    }
    ktdf.stage depends_in(%p0#3) depends_out(%p0#4) {        // compute stage holds a pipeline
      scf.for %m1 ... { scf.for %n1 ... {
        ktdf.pipeline {                                      // inner: vector granularity
          ktdf.stage ... { ktdf.data_transfer from %p0#0[..] ... to %fifo_in ... }   // L1 -> FIFO
          ktdf.stage ... { ... linalg.add ... }                                      // compute
          ktdf.stage ... { ktdf.data_transfer from %fifo_out ... to %p0#2[..] ... }  // FIFO -> L1
        }
      } }
    }
    ktdf.stage depends_in(%p0#4) depends_out(none) {         // bulk L1 -> DDR
      scf.for %m1 ... { scf.for %n1 ... { ktdf.data_transfer from %p0#2[..] ... to %C[..] ... } }
    }
  }
} }

第5步：探索循环空间

一旦确定管道的形状，剩下的问题是其中有多少部分会同时运行。在循环内部，一个阶段写入而另一个阶段读取的片上存储缓冲区会采用双缓冲机制，因此管道会在两个副本之间轮转（模2），DMA阶段可以在计算阶段仍在处理第N块时填充第N+1块。这正是解耦设计旨在提供的负载/计算重叠。随后，工作会被分配到每个核心内部的子核心实例（corelets）中，围绕核心本地管道的并行循环会变成 ktdf.parallel ... distribute(num_instances = N)，其中N从架构描述中获取。这为Inductor网格已分配的每个核心提供了第二层级的并行性。只有在完成这些步骤后，才会填充块大小的占位符，使得所选的大小已经考虑了并行实例的影响。

综合来看，这构成了一个真正的循环空间。循环结构、管道粒度、缓冲深度、实例分割和块大小都是IR中独立且明确的决策，每个决策都受到架构描述允许范围的约束。不过，我们需要明确当前的现状：在V1调度器中，这些选择是由本地启发式算法而非搜索或成本模型做出的。将所有选择以占位符形式保留在IR中，包括模因子和分布计数，是刻意为之的，以便后续可以添加成本模型或自动调优器而无需改动管道结构。

扩展一个管道：从简单开始，进行合法化，然后扩大规模

第6步：地址、单位和DataflowIR

最后一段将逻辑调度表转化为具体的内存布局。每个缓冲区在内存中获得一个地址，并与架构描述中的容量进行校验。快速存取存储器按程序分配，而设备内存分配则在不同程序间保持持久，因为这是程序间传递结果的方式。随后，抽象单元类型为每个核心（以及每个并行工作单元corelet）实例化，令牌依赖关系转化为显式信号，每个程序按引擎类型拆分为独立的数据流程序单元，仅保留在该引擎上执行的工作部分。在此过程中，数据传输转化为地址生成和发送/接收操作，计算单元ID变为每个核心的本地查询。最终每个程序形成两层结构：外层函数包含程序的真实签名并声明其符号，内层是无参数函数用于存储数据流中间表示。

符号值（无论是张量地址还是张量形状）在整个转换过程中保持完整。以地址为例，从内核参数计算出的起始地址会通过symbol.create_symbol生成符号ID，而生成该地址的表达式（参数加上该网格单元在张量中的份额，转换为单元地址粒度）则通过symbol.create_id在上一层记录。当不同网格单元从不同偏移量开始时，每个单元都会获得独立符号，通过单元级映射进行选择。

// 上一层：每个符号如何通过内核参数进行算术表达
func.func @local_schedule_0(%base: index) {
  symbol.create_id %base {symbol_id = -1}
  %w0 = arith.divsi %base, %c64 : index                   // 网格单元0，以64字节为单位
  symbol.create_id %w0 {symbol_id = -2}
  %b1 = arith.addi %base, %c49152 : index                 // 网格单元1：基地址+其内存块
  %w1 = arith.divsi %b1, %c64 : index
  symbol.create_id %w1 {symbol_id = -3}
  call @local_schedule_0_body() : () -> ()
}
// 数据流中间表示本身不接收参数，仅通过符号名称引用
%s0 = symbol.create_symbol {SymbolId = -2 : i64} : index
%s1 = symbol.create_symbol {SymbolId = -3 : i64} : index
%m  = uniform.def_immutable_mapping([%core0 -> %s0], [%core1 -> %s1]) : index
%a  = uniform.query_map(map: %m, key: %unit) : index
dataflow.get_logical_memory_view %ddr, %a ...

最终生成的符号表（每个符号都表示为内核输入的表达式）就是运行时在启动时最终要评估的内容。

数据流中间表示（DataflowIR）是关键交接点，它由数据流、矢量链、地址生成器、统一内存、符号系统和SCF方言构建的每个引擎程序单元组成。这是开源调度器的终止点，其对后续流程的唯一责任是生成包含稳定程序单元标签的有效DataflowIR，使用上述方言表示。

从DataflowIR到SpyreCode

剩余阶段依赖于具体设备实现，因此我们仅描述到与运行时部分连接所需的程度。

Symbolic SpyreBinary生成过程将每个程序的DataflowIR转换为SpyreBinary：用于在设备上配置和启动程序单元的二进制文件。任何依赖于调用方张量的值（无论是张量地址还是张量形状）都会以符号形式保留在二进制文件中。运行时符号替换的元数据生成过程会记录这些符号如何从内核的输入计算得出。最后，SpyreCode生成过程将SpyreBinary及其元数据与一个两阶段作业计划打包成两个部分。

SpyreCode包含的内容：符号化的SpyreBinary、其符号替换元数据以及一个两阶段作业计划

这两部分对应运行时后端在准备阶段和启动阶段的划分。准备阶段计划将SpyreBinary加载到设备上，且不依赖于调用参数，因此运行时的prepare_kernel只需执行一次。执行阶段计划会在主机上根据实际张量地址和形状计算符号值，将解析后的值传输到设备并启动计算，这对应运行时后端的三个启动阶段步骤。launch_jobplan会在每次调用时执行该计划，而那些在编译时已知所有地址和形状的内核只需执行启动操作即可。

尝试使用

Dataflow Scheduler基于上游LLVM/MLIR构建。README中给出了精确的LLVM版本号和构建方法。在完成MLIR构建后，调度器只需几个命令即可启动：

git clone --recursive https://github.com/torch-spyre/dataflow-scheduler.git
cd dataflow-scheduler
cmake -S . -B build -G Ninja -DMLIR_DIR=$MLIR_INSTALL_PREFIX/lib/cmake/mlir
cmake --build build

一个经典的"Hello World"示例是仓库测试目录中的简单逐元素加法操作：test/dataflow-scheduler-opt/ktir-to-dfir-sanity-check.mlir。这是一个小型的自包含KTIR内核，与本文前面提到的加法操作类似，跨两个计算单元运行，并针对上述示例设备进行目标设置。对它运行完整流水线会输出生成的DataflowIR：

./build/bin/dataflow-scheduler-opt -ktir-to-dfir \
    test/dataflow-scheduler-opt/ktir-to-dfir-sanity-check.mlir

添加--mlir-print-ir-after-all参数会显示每个阶段后的IR，这是观察程序拆分、三阶段流水线生成、路径扩展插入跳转以及流水线嵌套的最简便方式。对于自定义内核和设备，dataflow-scheduler驱动程序可以直接接收架构描述，使用方式为：dataflow-scheduler -kEmitDFIR --device=<path-to-device.mlir> <kernel.mlir>，并会将生成的DataflowIR写入文件。

结论

综合来看，这三篇文章描绘了从 torch.compile 到 Spyre 硅芯片的完整路径。Inductor 负责决定布局、分块、跨核心的工作分配，以及哪些张量保留在 scratchpad 中。KTIR 将这些决策以 MLIR 形式记录为计算块、显式内存、分布式视图、访问块和 Spyre 内置指令，但不涉及任何数据流细节。随后，数据流调度器将这些决策转化为实际执行，将内核拆分为程序，构建流水线，通过内存层次路由每个传输，并通过架构描述而非硬编码芯片知识对流水线进行嵌套、缓冲和并行化。设备特定阶段将结果转化为 SpyreCode，即一个符号化的 SpyreBinary 加上一个两阶段作业计划，其张量地址和形状可保持符号化直至启动，运行时只需准备一次并在每次调用时通过普通 PyTorch 流执行。

该路径上的每个边界都是一种中间表示（IR），每个 IR 都明确声明了下一阶段所需的事实：KTIR 中的放置信息、ScheduleIR 中的流水线、DataflowIR 中的引擎特定程序，以及 SpyreCode 中的准备/执行分离。将这些事实显式化，最终使数据流加速器能够通过 PyTorch 的普通接口运行。

数据流调度器版本 1 已经完成。它支持多核 KTIR 程序中带广播和约减维度的逐元素操作，并已在代表性 Transformer 操作上完成了从 KTIR 到调度和代码生成的端到端验证。迄今为止最显著的收益之一是开发人员引入新操作所需的时间大幅减少：我们启用的第一个操作比之前架构快约 4 倍，后续操作则快约 8 倍。同样的方法也使我们从 Inductor 和 Triton 前端都获得了可行的路径。接下来，我们将完善数据流调度器中的矩阵乘法支持，并添加基于成本模型的性能优化，目标是实现 LLM 的端到端运行。

KTIR 和数据流调度器已开源，上方的“Try it”部分是最快查看流水线运行效果的方式。如果你正在构建面向数据流或 scratchpad 加速器的编译器，我们非常希望听到你的反馈，因为架构描述方法的设计初衷是让调度器在 Spyre 之外的场景中同样具有实用价值。

/post-content

/inner-wrap