TL 证明，编译器驱动的映射可以在空间数据流加速器上提供厂商级别的性能。在 GEMM 和 FlashAttention 上，以最少的人工干预实现了与厂商库足以竞争甚至更好的性能，将手工设计的内核转变为可重用编译流水线的输出。同时，TL 的硬件抽象使得编译器过程在很大程度上可在不同加速器间重用，简化了对多样化空间数据流架构的支持，并为未来芯片的设计空间探索提供了自然的基底。
关键词：空间数据流架构、编译器框架、基于 tile 的程序、数据重用、硬件表示、性能建模
面向新型芯片架构的编译器新范式，如何将瓦片式程序自动映射到空间数据流加速器上。
空间数据流架构正成为下一代高性能计算系统的关键方向。与传统的冯·诺依曼架构（如 CPU 和 GPU）不同，它通过芯片上高带宽的片上网络在计算单元间直接传递数据，从而大幅降低了对高延迟、低带宽的全局共享内存的依赖，显著提升了能效与计算密度。
这类架构的典型代表包括 Tenstorrent、Cerebras、Graphcore、SambaNova、Groq 等公司的芯片，它们将大量处理核心组织成二维或更高维的网格，每个核心具备本地内存，并通过可编程的网络直接通信。
然而，空间数据流架构的极高性能潜力和效率，伴随的是其 “编程难” 的核心挑战。性能极度依赖于工作负载如何映射到硬件上：哪个核心执行哪部分计算、数据如何在本地内存间划分、通信如何调度。
一个不良映射会导致负载不均衡、网络拥塞、核心利用率低下或片外访存过多，性能甚至会严重劣于传统 GPU。因此，大多数用户只能依赖厂商为少数热门算子（如矩阵乘、卷积、注意力）提供的手工优化库，限制了硬件潜力的发挥与新算子的快速部署。
- TL: Automatic End-to-End Compiler of Tile-Based Languages for Spatial Dataflow Architectures
- https://arxiv.org/pdf/2512.22168
今天要介绍的这项工作 —— TL（TileLoom） ，提出了一套完整的端到端编译器框架，能够将基于瓦片的高级程序（如 Triton 编写的内核）自动编译到空间数据流架构上。
其核心创新在于：将传统上由硬件运行时或手工库完成的跨核心映射、数据移动调度与优化，全部提升到编译时完成，并支持多种不同的空间架构。
TL 框架基于 MLIR 生态构建，前端将 Tile 级内核转换为标准中间表示，数据流规划阶段确定时空映射与数据移动方案，后端生成硬件特定可执行文件。各模块在架构表示与性能模型指导下协同工作，解决传统编译器仅优化单 Tile 代码生成的局限，实现 Tile 实例在空间分布式核心上的高效分配，提升数据复用率并降低通信开销。
本文实验中，TL 在 Tenstorrent Wormhole 系统上，针对 GEMM 和 FlashAttention 等核心算子，性能平均达到了厂商提供库的 1.03 倍和 1.91 倍，且在一些场景下优势显著。
一、空间数据流架构的机遇与编程之困
1.1 DRAM 高能耗、SRAM 小且贵
现代高性能负载，尤其是深度学习负载，是高度数据密集型的。对多数负载而言，瓶颈已不在计算，而在内存带宽。
随着工艺演进，芯片可容纳的运算单元数量持续增长，但片外内存的带宽与容量并未同步提升。每次 DRAM 访问的能耗远高于一次算术运算，片上 SRAM 也随制程缩小变得更加昂贵。这些趋势使得以内存为中心的冯·诺依曼架构难以充分利用其计算单元。
1.2 空间数据流架构的优势：高效数据访问
空间数据流架构提供了不同的思路：它们将计算组织在显式的、通常由软件控制的数据移动之上，数据通过片上网络和缓冲区在核心间传递，减少了对大型共享缓存和高延迟片外内存的依赖。
当数据通过短距离的片上连线在核心间传递时，每比特能耗和延迟远低于往返共享缓存或片外内存。
以 Tenstorrent 的系统为例：其核心被组织为二维网格，每个核心通常为 SIMD 或向量引擎（常含矩阵单元），并配备本地暂存器，通过分组交换的 Mesh 片上网络连接。
64 个核心可并发访问其本地暂存器，实现约 24.5 TB/s 的总带宽，远高于 NVIDIA H100 的 6 TB/s L2 带宽。如此高的核心级带宽显著缓解了带宽受限算子的瓶颈，使内核能维持更高的吞吐量。
1.3 空间数据流架构的挑战：可编程性差
然而，空间数据流架构在提供高效能潜力的同时，也带来了严重的可编程性问题。性能高度依赖于负载映射方式：哪些核心执行计算的哪些部分、数据如何在本地内存间分区、流量如何在网络上调度。
一个朴素的映射可能导致严重的负载不均衡、网络拥塞、核心利用率低下或过多片外访存，通常导致性能和能效极差。
1.4 常见编程模型抽象：CUDA block-thread 模型
一种常见的并行加速器抽象是 CUDA 推广的 grid-block-thread 模型。
* 程序员将问题分解为多个块，并在设备上启动由块实例组成的网格；
* 每个块（如 CUDA 线程块或 OpenCL 工作组）对输入或输出张量的不同区域执行相同计算，网格共同覆盖整个问题域。
类似的基于网格的瓦片模型也出现在 Halide 和 TVM 等高级系统中，它们将瓦片形状和启动配置视为调度参数，并自动搜索以改善局部性和并行性。
在 GPU 等传统以内存为中心的架构上，网格级别的管理由硬件完成：硬件调度器动态地将块分配给流式多处理器，共享缓存层次结构隐式捕获块间的大部分数据重用，无需依赖显式的块间通信。这使得编译器对网格级布局、执行顺序或通信模式的控制空间有限。
相比之下，块-线程级别是软件定义的：编译器或程序员决定块内工作如何映射到线程、线程束和 SM 内资源。CUDA 在此级别暴露细粒度控制，但充分利用需要仔细调优并对 GPU 微架构有深入了解。
为使块级（瓦片级）编程更易用，出现了 Triton、TileLang、CuTile、Tilus 和 Taichi 等语言，它们让用户用高级瓦片操作符表达单个块的工作，编译器将这些操作符下译到核心内或 SM 内资源，并处理相关的微架构陷阱。
1.5 空间数据流架构：将控制权交给软件的可编程性挑战
空间数据流架构改变了这一局面。这些架构将计算和暂存器分布在由高带宽片上网络连接的大型处理单元阵列中，通常不具备能自动利用跨核心重用的大型统一缓存。它们将通信和布局决策的负担从硬件缓存和调度器转移到了软件，暴露了更丰富的跨核心可编程性，让软件显式决定核心间如何通信。
因此，要获得良好性能，编译器或程序员不仅需要知道如何在单个核心上实现瓦片程序，还需要知道如何将瓦片实例分布在核心间，以及如何在时间上调度它们，以便数据能在片上网络上被高效转发或多播，并有效利用内存系统。
额外的跨核心可编程性创造了巨大的映射设计空间：有效映射的组合空间极为庞大，不同映射暴露不同的重用模式和通信成本。在当前系统中，这些映射决策被固化在特定厂商的编译器和库中，这些库编码了特定架构的布局、路由和流水线策略，但开发耗时，且不易泛化到新内核、模型或硬件代际。
通过给予软件更多控制权，空间数据流架构更难编程，但可以简化硬件并提高效率。关键在于，随着硬件通过暴露显式的核心、网络和内存而非不透明的缓存和调度器，使其对软件更加透明，数据流架构比以内存为中心的架构更具可预测性。
在掌握足够硬件信息的前提下，编译器更容易推理不同调度方案的性能，并推导出良好的静态映射。基于这一洞见，我们提出了 TL。
二、TL 编译器框架总览：自动化的端到端瓦片映射
TL 编译器框架旨在将用基于瓦片的 DSL 编写的内核（如 Triton）编译成目标空间数据流架构上的可执行文件。其编译栈围绕三个主要组件构建：
* 一个前端，用于将瓦片级内核下译为一种与数据流无关的标准 MLIR 表示；
* 一个数据流规划阶段，用于确定时空映射、数据移动，并在一种数据流感知的标准 MLIR 表示中生成候选方案；
* 一个后端，用于为每个核心生成特定硬件的可执行文件。整个过程由多层次架构表示和性能模型指导。
2.1 前端
本阶段接收瓦片级内核以及关于该内核如何在整个问题上（启动网格）进行扩展的描述。它探索候选的块形状：瓦片大小和布局，并构造相应的程序。这些候选程序随后被下译为基于 MLIR 的中间表示，并通过规范化处理，使其共享适合数据流规划流水线的通用结构。
此时，计算被表示为统一、与数据流无关的中间表示（IR）：它编码了瓦片和网格结构，但尚未承诺任何特定的硬件映射。
2.2 数据流规划阶段
本阶段决定这个逻辑瓦片网格如何在目标架构上实现。在硬件描述的指导下，TL 探索瓦片实例在核心和时间上的时空映射。
* 对于每个候选映射，它执行数据重用分析，以识别哪些瓦片可以通过片上网络共享数据，或通过缓冲随时间重用；
* 并推导出具体的数据移动计划：每个张量瓦片分配在哪里、何时复制、以及如何通过缓冲区广播或共享。
时空映射和数据移动计划共同生成了一个潜在最优数据流规划候选的设计空间。
2.3 本文工作
为了建模任意数据流架构，本文提出了 架构表示，它为数据流规划搜索和性能模型提供关键输入。
* 扩展描述（跨核心）捕捉核心阵列和互连的空间结构，并指导时空映射决策。
* 核心内描述捕捉本地内存层次结构和计算资源，并指导关于数据存储位置和暂存方式的决策。TL 结合这两个抽象层次来生成性能模型，该模型估计不同数据移动计划的成本，考虑内存带宽、片上网络能力和每个核心的计算吞吐量。性能模型用于从数据流设计空间中选择最优候选。
在数据流规划阶段之后，中间表示变为数据流感知的：内存分配、复制和通信端点根据所选映射具体化。从设计空间中选择最优候选后，块级优化和代码生成阶段将每个核心上运行的数据流感知程序编译到厂商现有的后端，为每个核心生成可执行文件。这对应于现有的将基于瓦片的 DSL 编译到块级程序的编译器。
这些组件共同使 TL 能在不同程度的硬件细节下支持一系列空间数据流架构。
TL 采用两步选择策略：自动生成的性能模型首先对候选数据流计划进行排序并静态选择前 k 个映射，然后这 k 个候选在实际硬件上进行性能剖析以选择最终映射。
这种模型指导搜索与硬件验证相结合的方式，使 TL 能生成 高质量的映射，同时保持在 不同空间数据流系统间的可移植性。
三、TL 核心机制深入解析
3.1 时空映射：逻辑网格到物理核心的映射艺术
在数据流规划流水线之前，需要一个统一的、与数据流无关的 MLIR 表示。假设此前端已生成此表示，本文将在后续章节介绍 Triton 前端的实现细节。清单 1 展示了一个用于矩阵乘法内核的此类 MLIR 示例。块大小以及输入和输出矩阵的步幅由前端固定。
二维输出空间在两个网格维度 x 和 y 上进行分区，这些维度在一个 m×k×n 矩阵乘法（一种输出固定的瓦片化）的 m 和 n 维度上进行扩展。跨瓦片的扩展由在 block_id_x 和 block_id_y 上的仿射并行循环表示。在此循环内部，一个 scf.for 循环在 k 维度上迭代并累积到相同的输出瓦片中，表示在一个输出瓦片上单个块内的顺序执行。
前端需要对内存操作的地址运算进行“仿射化”，因此每个加载和存储地址都是瓦片索引和瓦片内索引的仿射函数。瓦片级计算本身用 linalg 操作表示；程序的这部分在数据流规划期间保持不变，稍后由后端下译。
时空映射的目的是决定仿射并行循环的迭代空间——即逻辑多维瓦片网格——如何分配给物理核心和时间。为了保持局部性，TL 使用基于瓦片的映射：迭代空间的连续区域被映射到核心阵列的连续空间区域或执行调度中的连续时间区域。
在如图 1 所示的二维网格架构上，时空映射产生清单 2 中所示的循环结构。
最外层的仿射并行循环现在在硬件空间维度 x 和 y 上迭代，每个维度大小为 8。这些索引直接对应于 8×8 二维网格中的核心。映射后，此循环表示实际上在核心间并行运行的代码；其语义已从可并行化工作项的逻辑网格转变为物理并行核心索引。
接下来的仿射 for 循环在 tx 和 ty 上迭代。这些循环枚举分配给相同硬件阵列的瓦片波次，即跨块的时间维度。每个波次将一批逻辑瓦片分配给可用的空间核心，这些波次的顺序决定了瓦片遍历阵列的时间调度。最内层的 scf.for 在每个核心的程序内仍然是一个纯粹的顺序循环。
设计空间 。从原始并行维度到空间和时间维度的映射定义了内核的数据流。在我们基于瓦片的方案下，设计空间由三个相互耦合的选择来刻画：
- 首先，每个原始并行维度可以映射到零个或多个空间维度 。将一个并行维度映射到一个空间维度对应于按该空间维度的大小对该循环进行瓦片化，并引入一个新的、覆盖硬件索引的外部仿射并行循环。
- 其次，当一个并行维度被多个空间维度瓦片化时，瓦片化的应用顺序很重要 。不同的瓦片化顺序在网格上诱导出不同的瓦片空间布局和不同的执行调度，因此暴露出不同的空间重用机会和不同的通信成本。
- 第三，一旦所有可用的空间维度都被使用，剩余的并行维度就变成作为波次循环实现的时间维度 ，如清单 2 中的
tx和ty循环。 这些时间循环及其顺序决定了瓦片如何随时间分批加载到阵列上，以及它们以何种顺序重新访问全局迭代空间的每个区域，这反过来影响时间重用和片上网络流量的形态。
TL 通过探索这些选择的组合来枚举候选的时空映射。每个映射固定了一个具体的循环嵌套结构，后续的分析使用该结构来推理数据布局、重用和通信。
3.2 数据重用分析与内存操作映射
不同的时空映射暴露出不同的跨时间和跨核心的数据重用机会。 TL 首先分析这些机会，然后决定如何将数据分配到内存中，以及何时何地通过片上网络广播或从全局内存加载来发起复制操作。
仿射访问的重用分析
对于一个固定的时空映射，循环嵌套包含：覆盖硬件核心索引的空间循环（仿射并行）、覆盖瓦片波次的时间循环（仿射 for）以及每个核心内部的顺序循环（scf.for）。
前端将所有内存访问表达为这些循环索引的仿射函数。对于每次访问，TL 检查其仿射表达式中出现了哪些归纳变量。如果一次访问不依赖于诸如 sx 之类的空间索引，则沿该维度所有核心访问的瓦片是相同的，并在该维度上空间可重用。
- 如果一次访问不依赖于诸如
tx之类的时间循环变量，则同一个瓦片在该时间循环的所有迭代中使用，并在该维度上时间可重用。 - 如果访问仅依赖于顺序索引，则重用纯粹是核心内的。TL 将这些信息记录为内存操作上的重用注解。
空间重用与广播
从一个保守的基线开始：在此基线中，每个核心在最内层循环中直接从全局内存加载其瓦片，核心间没有显式共享。
如果一个加载操作没有空间重用，则该瓦片对每个核心是唯一的，因此加载必须保持为每个核心的全局内存操作。如果一个加载操作沿一个或多个空间维度空间可重用，TL 可以用更少数量的全局加载（随后通过片上网络进行广播）来取代许多每个核心的加载，从而减少全局流量。
在最简单的情况下，如果一个瓦片仅沿单个空间维度可重用，一个指定的生产者核心（或一小组生产者）从全局内存加载一次，并沿该维度转发，而接收核心则缓冲其本地副本。
当一个瓦片沿多个空间维度可重用时，有几种具体方法来实现该重用。
- 一种选择是先在所有行（或列）上复制瓦片，然后沿并行维度执行独立的一维广播；
- 另一种是以波前式模式在阵列上传播瓦片。这些选择在片上网络流量、延迟和本地缓冲区使用之间呈现出不同的权衡。
清单 3 展示了矩阵乘法映射到二维网格示例架构的一个候选示例（即图 1），采用二维数据流，其中 A 瓦片通过片上网络水平链路为每一行核心广播，B 瓦片通过垂直链路为每一列核心广播。广播信息作为注解与加载指令关联 。
TL 并不固定单一策略：它使用硬件表示中的网络描述来枚举合法的广播模式，对于每个空间可重用的加载操作，生成一小组候选实现，范围从直接每核心全局加载到一维和二维广播。TL 枚举所有内存操作所有可能的组合，从而创建一个设计空间。它们不同的硬件成本稍后由性能模型考虑，以选择最佳方案。
时间重用与循环提升
时间重用通过选择加载操作的发起循环级别来实现 。此时，循环顺序已由时空映射固定；但我们可以将加载向外提升，以便同一瓦片在更多迭代中被重用，代价是在本地缓冲区中保留更长时间。
考虑一个简化的类 GEMM 循环嵌套，假设一个候选循环顺序为 tm → tn → tk，访问 A[tm, tk] 依赖于 tm 和 tk，但不依赖于 tn，因此 A 的瓦片在 tn 循环上时间可重用。如果我们将 A 的加载提升到 tk 循环之外，我们必须为当前 tm 缓冲所有瓦片 A[tm, *]。由于地址依赖于 tk，跨 tk 提升将缓冲区域从单个瓦片 A[tm, tk] 扩大到整个条带 A[tm, 0..K_tiles-1]。进一步向外提升，超过 tn，保持相同的缓冲条带，但在 tn 的所有值上重用该条带。
通常，将一个加载提升到其访问不依赖的循环上会增加重用而不增加缓冲区域大小，因为对于该循环的所有迭代，访问的瓦片是相同的。将一个加载提升到其访问依赖的循环上会按该循环的范围成比例地扩大缓冲区域，因为必须同时保持更多不同的瓦片存活。
TL 应用这些规则为每个加载或广播枚举所有合法的提升级别。对于每个级别，它计算所需的缓冲区足迹，并丢弃那些足迹超过硬件模型容量的选项。
时间与空间重用
时间重用和空间重用是正交的。 一个瓦片可以仅时间重用（每个核心加载一次，跨迭代重用）、仅空间重用（广播一次并立即消耗），或以两种方式重用（广播一次，然后在多个时间迭代中重用） 。在所有情况下，该决策可视为在固定循环顺序下，选择瓦片首次加载或接收的时间点，以及它在本地存储中保持存活的时间长度。
为所有加载操作组合这些选择，产生一个具体的分配和复制映射：描述每个瓦片在每个时间点驻留在哪个内存中，以及发生哪些片上网络传输。此调度可用带注解的循环结构表示，清单 5 展示了一个简化示例。
每个内存加载都注解了目标缓冲区、加载类型以及所需的片上网络资源。TL 修剪违反内存容量约束的映射，并将剩余的候选传递给性能模型，该模型评估它们的计算、内存和网络成本，并选择前 k 个映射用于后端剖析。
3.3 硬件表示：统一描述异构架构
TL 旨在面向多种数据流架构。为了做出映射决策，编译器需要一个结构化的硬件描述：核心如何在空间中排列、内存放置在何处、组件如何连接以及每个位置有哪些计算资源可用。
TL 将这些信息捕捉在一个多层硬件表示栈中。此栈的不同层次由编译器过程的不同阶段消费。
本文将此表示编码在一个自定义的 MLIR 方言 df 中。该方言提供了描述机器扩展结构（核心和互连）、内存层次结构及其连接性以及核心内计算单元的操作符。性能模型和映射过程一起处理用此方言编写的程序和硬件描述，而非硬编码任何特定架构。
扩展架构
在最上层，df 方言描述核心的空间布局和片上互连。需要以下操作符：
df.spatial_dim(size)：声明一个抽象空间维度，用于索引和复制硬件组件。空间维度自然表示并行资源（如核心或内存）的数组。df.core(scaleout, scalein)：声明一组由scaleout中列出的维度索引的核心；该操作的scalein参数包含位于每个核心内部的计算组件（可选参数，用于较低抽象层次）。df.interconnects(components, map, bandwidth)：声明一个网络（一组链路），根据仿射映射连接一组组件，bandwidth指定每链路的带宽。
这些操作符用于描述时空映射中使用的扩展架构。对于二维网格示例架构，可以如下描述：
%x = df.spatial_dim 8
%y = df.spatial_dim 8
%cores = df.core { scaleout = (%x, %y) }
%noc_h = df.interconnects %cores, %cores { map = affine_map<(d0, d1) -> ((d0 + 1) mod 8, d1)>, bandwidth = 28 }
%noc_v = df.interconnects %cores, %cores { map = affine_map<(d0, d1) -> (d0, (d1 + 1) mod 8)>, bandwidth = 28 }
此片段描述了一个 8×8 核心阵列，通过水平和垂直环（或环面）连接。仿射映射中的模运算编码了环绕链路。时空映射过程使用此扩展描述将逻辑瓦片网格映射到物理核心，性能模型使用 df.interconnects 操作符来估计内存操作的通信成本和流量拥塞。
内存与数据移动
在规划数据移动时，TL 必须推理具体的缓冲区以及网络如何将数据馈送到其中。为此，扩展描述需要用显式的内存进行细化。
df.memory(scaleout, size, bandwidth)：声明一组由 scaleout 维度索引的内存。每个实例具有给定的容量和每端口带宽。df.mux(dst, bandwidth, srcs, map)：声明 dst 组件与 srcs 之间的一对多连接，拓扑由仿射映射指定。该操作符捕捉诸如“每个核心可访问其本地暂存器”或“核心组共享一个 DRAM 通道”之类的扇出连接。
包含 L1 内存和 DRAM 的二维网格描述的一个降级版本如清单 7 所示。
此表示现在区分可以容纳瓦片的物理缓冲区。例如，它指定每个核心 (x, y) 有一个约 1.5 MB 的 L1 暂存器，并且核心间流量在 L1 之间流动，而非直接在核心之间流动。它还编码了 DRAM 通道的连接方式：每个边缘上的每组四个相邻核心共享同一个 DRAM 存储体。当 TL 选择在何处缓冲瓦片以及内存操作的成本时，它使用这些 df.memory、df.mux 和 df.interconnects 操作符。如清单 5 所示，加载指令被注解为与这些物理资源的绑定。
核心内计算模型
为了将性能模型驱动到单个核心的级别，TL 需要对每个核心内的微架构进行粗略描述。df 方言为此目的提供了操作符：
df.mat(shape, throughput)：声明一个具有给定输入形状和持续吞吐量的矩阵单元（例如张量核心）。df.vec(shape, throughput)：声明一个具有给定向量宽度和吞吐量的向量单元。df.scalar(latency)：声明一个标量单元，具有给定标量操作的延迟。
假设每个单元接受指定形状的操作数，并以给定的吞吐量产生结果。然后这些单元通过 df.core 的 scalein 参数附加到核心，该参数描述每个核心的内部组成。因此，二维网格描述的最低层级版本通过核心内单元扩展了之前的清单：
%FPU = df.mat {shape=[32, 32, 32], throughput=98}
%SFPU = df.vec {shape=[32], throughput=3}
%cores = df.core {scaleout = (%x, %y), scalein=(%FPU, %SFPU, [8, 1])}
有了这些信息，性能模型可以推理计算和内存的时序。例如，它可以估计特定的瓦片级矩阵乘法在矩阵单元上消耗多少周期，以及片上网络带宽和 L1 缓冲区是否足以持续供给数据。
超越二维网格的表达能力
尽管我们使用了 Tenstorrent 二维网格架构作为示例，但 df 方言可以描述其他空间数据流架构。例如，一个类似于 IBM-Spyre 加速器的一维三环拓扑可以用 df 程序描述。
讨论
我们将硬件表示构建为抽象层栈，因为每个编译器过程应该只依赖于它实际需要的细节层次。时空映射只需要核心的扩展结构以及互连的拓扑和带宽。数据移动规划还需要知道内存放置在何处以及它们如何连接到计算和 DRAM。细粒度性能模型反过来需要核心内计算单元及其吞吐量的高层视图。
这种分离提高了编译器的可重用性。更改片上网络、内存层次结构或每核心微架构只需修改 df 描述，而无需重写优化过程以及硬件描述的其余部分。它还搭建了从软件级映射决策到硬件级设计权衡的桥梁。从最低抽象层级开始，相同的表示可以进一步细化，以包含实现特定的成本如面积和功耗，实现对映射和硬件配置的联合设计空间探索——这是空间数据流架构的一项重要能力，因为其架构本身在不同代际和厂商之间存在很大差异。
3.4 性能建模：从抽象映射到时间预估
在时空映射和数据重用/分配决策之后，TL 拥有一组候选数据流调度方案。每个候选表示为一个 MLIR 程序，其中循环嵌套、内存操作和数据移动（全局加载、广播、缓冲加载）都已完全指定，并且每个内存或网络操作都绑定到 df 方言中描述的具体硬件资源。清单 5 展示了一个简化示例。
性能模型的作用是使用 df 描述中定义的计算单元、内存和互连来估计每个候选的执行时间，然后为下游代码生成和剖析选择前 k 个候选。
图 4 说明了性能模型如何评估清单 5 示例的整体执行时间。它从最内层循环向外评估每个候选，分层聚合计算、内存和网络成本。
每个循环体的计算成本
我们首先估计最内层循环体的执行时间，将其视为在单个核心上运行的块级程序。对于给定的瓦片形状，每个高级操作符被分解为核心的低级计算固有指令。可用的矩阵、向量和标量单元及其吞吐量来自通过 df.core 附加到该核心的 df.mat、df.vec 和 df.scalar 单元。
对于每个操作符，TL 使用其 linalg 语义来恢复并行迭代空间。这告诉我们，对于每种功能单元类型，该类型有多少固有调用是独立的，并且原则上可以并行发起。然后，我们概念上将此固有指令调度到可用的同类型并行单元上：如果有 N 个独立实例映射到具有 U 个相同单元的功能单元类型，每个单元每个周期能发起 r 个固有指令，我们近似该操作符在该功能单元类型上的时间贡献为 ceil(N / (U * r)) 个周期。
然后，TL 考虑不同功能单元类型之间的数据依赖性和资源共享。针对不同功能单元类型的独立操作符可以并行执行，而相互依赖的操作符或竞争相同功能单元类型的操作符必须顺序执行。循环体计算时间近似为顺序段的时间总和，其中每个段的时间是该段内可以并行运行的所有操作符的最大值。
此模型暴露的潜在并行性在具体的微架构上不一定能完全实现，但模型并不试图精确模拟核心的指令调度器。相反，它经过校准，足以准确区分计算受限和内存受限的映射，并推理计算与数据移动之间的重叠。在我们的实验中，这种粗粒度计算建模足以区分不同的数据流调度方案。
计算-内存重叠
一旦我们有了循环体的估计，我们纳入数据移动。设 T_load 为循环体一次迭代中所有加载花费的时间，T_comp 为计算时间，T_store 为所有存储花费的时间。我们假设每次迭代作为流水化的加载-计算-存储序列执行，使用双缓冲：当迭代 i 在计算时，迭代 i-1 的存储和迭代 i+1 的加载在可能的情况下并行进行。
对于一个有 I 次迭代的最内层循环，总执行时间近似为：
T_total ≈ max(T_load, T_comp, T_store) * (I - 2) + T_load + T_comp + T_store
第一项对应 I-2 个稳态迭代，其中加载、计算和存储可以重叠，吞吐量受限于 T_load、T_comp 和 T_store 中的较慢者。剩余项对应流水线的填充和排空。此行为如图 4 所示，其中 k 循环体与加载下一个 B 瓦片以及存储前一次迭代的 C 结果瓦片并行执行。
并发数据传输与网络流量
多个内存操作可能同时发生，并在片上网络上产生流量争用。TL 使用硬件表示中描述的互连结构以及内存操作映射期间附加的映射注解，估计在此争用下每个内存操作的有效带宽。在映射步骤中，每个加载或存储被降级为全局加载/存储或广播模式，编译器记录它使用了网络的哪些子集。
- 对于全局加载，我们假设访问足够随机，流量分布在片上网络链路上。
- 对于广播，使用的资源取决于选择的模式：在二维网格示例中，沿每行独立执行的广播仅使用水平环链路，而在整个网格上的广播可能同时使用水平和垂直环。这些选择在内存操作映射期间固定，并作为注解出现。
给定这些注解，性能模型根据内存操作占用的网络链路和内存接口对其进行分组。对于共享特定链路子集的每组操作，它聚合它们提供的流量，并通过在它们之间划分标称链路带宽来推导每次传输的有效带宽。等效地，任一操作看到的带宽按使用相同链路或存储体的并发传输数量成比例减少。
然后，每次加载或存储的传输时间根据其瓦片大小和有效带宽计算。每次迭代的加载和存储时间 T_load 和 T_store 通过跨操作组合这些传输时间获得，将在不相交链路集上的传输视为并行运行，在重叠链路集上的传输视为分时共享相同资源。最后，将这些 T_load 和 T_store 值代入上述流水化重叠模型。
候选排序
对于每个候选调度，TL 结合计算和数据移动估计来获得近似的端到端执行时间。该时间反映了计算与通信的平衡、空间和时间重用的好处，以及片上网络和内存争用对 df 描述的具体硬件的影响。
TL 然后按此估计时间对所有候选排序，并仅保留前 k 个数据流映射。只有这些前 k 个候选被交给后端进行完整的代码生成和硬件上剖析，其中选择最终性能最佳的配置（Top1）。k 的选择控制着编译时间与包含真正最优解可能性之间的权衡：更大的 k 探索更多映射，但需要更多编译和剖析时间。我们在后续章节研究这一权衡。
四、实验评估：性能如何？
实验设置：所有实验均在连接至主机的 Tenstorrent Wormhole n300d 卡上进行。同一 CPU 平台用于运行整个 TL 编译器栈。Tenstorrent Wormhole n300d 是一个空间数据流加速器，由两个插座组成。每个插座集成 64 个 Tensix 核心，运行在 1 GHz。每个 Tensix 核心每个周期最多支持 1024 次 FP16 操作，每个插座的峰值吞吐量为 64 TFLOp/s。每个插座提供 96 MB 的片上 SRAM 和 12 GB 的 GDDR6 DRAM，峰值带宽为 288 GB/s。
所有实验使用相同的物理 Wormhole n300d 设备，但评估三种核心阵列的逻辑配置：完整的 8×8 对称二维网格、作为非对称矩形网格使用的顶部 4×8 子网格，以及作为一维环使用的单个 1×8 行。这些配置近似于不同的扩展拓扑，并允许我们研究空间结构形状如何影响最优数据流映射。
使用 df 方言对底层 Tenstorrent 架构进行建模。由于无法获取完整的专有硬件规范，我们通过剖析恢复关键参数：矩阵和向量单元的吞吐量，以及片上网络和 DRAM 通道的有效带宽，使用独立的微基准进行测量，并插入到 TL 性能模型使用的硬件表示中。
评估 TL 在两个具有代表性且对性能关键的算子：通用矩阵乘法和 FlashAttention。对于这两种工作负载，在广泛的输入形状范围内对 TL 进行基准测试，以全面评估其稳健性和性能特征。
虽然 TL 不绑定到特定的前端语言，但其实现在本文中针对广泛使用的基于瓦片的 DSL Triton。我们调整 Python 中的瓦片形状，与 Triton 内核一起，并使用 triton-shared 为每种配置生成 MLIR。然后运行自定义的仿射化过程，将所有内存访问的索引运算重写为仿射表达式，并运行其他规范化过程以标准化输入 MLIR 格式。产生的 MLIR（作为 TL 数据流优化堆栈的输入）具有清单 1 所示的形式。
后端：我们将数据流感知的 MLIR 下译到 Tenstorrent 的低级 C API TT-Metallium，以生成最终可执行文件。TT-Metallium 暴露粗粒度的计算和数据移动固有指令，并执行大多数块/核心级优化。TL 将指令选择和细粒度核心内调度委托给 TT-Metallium 后端。
4.1 GEMM 性能评估
首先评估通用矩阵乘法。对于每个问题形状和硬件配置，TL 使用性能模型选择前 5 个数据流候选，为它们生成代码，在硬件上剖析所有五个，并报告性能最佳的一个（Top1）。与 Tenstorrent 的厂商库 TTNN 进行比较，TTNN 在其两种手工编写的数据流模板 TT-1D 和 TT-2D 之间根据其自定义策略和输入形状进行选择。此选择决定了空间重用策略，并与其他超参数正交，例如块大小，这些超参数使用其他自定义策略决定。
- TT-1D 使用一维广播模式：对于每个 GEMM，两个输入矩阵中较小的一个由每个核心从全局内存加载，而另一个输入在整个阵列上广播。
- TT-2D 使用二维广播模式：两个输入矩阵在网格上流动，一个从顶部，一个从左侧，使瓦片以脉动式波前传播。TTNN 使用 Tenstorrent 自己的策略在 TT-1D 和 TT-2D 之间选择。TTNN 还根据输入形状确定单个块大小。
相比之下，TL 搜索更大的映射空间。它不仅探索一维和二维广播模式，还探索将逻辑维度分布到空间维度的不同方式、循环嵌套中内存操作的不同提升点（时间重用）以及不同的块大小。
图 5 报告了三种扩展架构上 140 种 GEMM 配置的性能。
- 第一行对应 1×8 环，第二行对应 4×8 非对称网格，第三行对应完整的 Tenstorrent 单芯片 8×8 对称网格。
- 在图的每一列中，对于形式为 C[M,N] = A[M,K] × B[K,N] 的 GEMM，我们固定 K，并将 M 和 N 从 256 扫描到 16384，如 x 轴所示。
- 所有结果相对于相同配置下 TTNN 的性能进行归一化（越高越好）。在标准的 8×8 网格上，TL 相比 TTNN 实现了 2.8% 的几何平均加速比，在所有配置的 78.5% 上性能差异在 10% 以内，在 62.8% 的情况下优于 TTNN，加速比最高达 2.19 倍。
TL 在非常小的形状上往往表现不佳，在这些情况下绝对运行时间由我们的模型未捕获的硬件开销主导，并且由于缺乏专有硬件细节而难以纳入。这一点也可以从后续关于性能模型的讨论中观察到，如图 9 所示，在小形状区域预测精度下降，这限制了性能模型选择最佳映射的能力。
相对于固定的 TT-1D 和 TT-2D 模板，TL 分别实现了 30% 和 9% 更高的几何平均性能，突显了选择正确数据流的重要性。固定的数据流策略在不同的形状和硬件拓扑中通常不是最优的。
此外，大多数配置在 (M,N=16384,1024) 处有一个共同的峰值，其中 TL 和 TT-1D 显著优于 TTNN 基线，表明 TL 正确选择了类似于 TT-1D 的调度，而 TTNN 由于其固定的调度策略未能做到这一点。
数据流规划也与网格配置和算子形状密切相关。
例如，在图 5 的第 5 列中，TT-1D 和 TT-2D 的相对性能在 4×8 非对称网格与 8×8 对称网格之间发生了逆转：TT-1D 在非对称网格上更好，而 TT-2D 在对称网格上更好。TTNN 和 TL 都使它们的数据流决策适应这一变化，但 TL 更一致地选择与实际性能最匹配的映射，因此实现了更高的整体吞吐量。
4.1 不规则形状的 GEMM
为了进一步说明数据形状如何影响数据流映射的选择，我们评估了两个不规则 GEMM 形状族。第一种设置固定 M = N = 32768，K 从 256 变化到 2048。第二种设置固定 M = K = 32768，N 在相同范围内变化。图 6 总结了结果。
- 当变化 K 时，TL 和 TTNN 都表现出与一维和二维基线类似的行为。这是因为 K 维度在每个核心内顺序映射，主要影响核心内计算成本，留给数据流优化的空间相对较小。
- 当变化 N 时，首选数据流会发生显著变化。
- 随着 N 增加并接近 M，工作负载在行和列维度之间变得更加平衡，类二维的广播策略变得更具吸引力，因为它们可以沿网格的两个维度高效地重用数据。
- 当 N 远小于 M 时，类一维的策略更有效，这一趋势也体现在 TT-1D 和 TT-2D 的相对性能中。从一维有利区域到二维有利区域的确切转换点取决于块大小、计算内存比以及其他硬件特定因素。这种敏感性突显了 TL 成本模型指导搜索的必要性，它显式地考虑了空间重用、通信量和架构约束，而不是像 TTNN 那样依赖固定的、手动设计的启发式方法。例如，在图 6b 中 N = 1024 处，TTNN 选择了一个类似于二维基线的映射，并且由于其固定的选择策略，未能识别其自身的一维模板在此形状下性能更优。
4.2 FlashAttention 性能评估
接下来，我们评估另一个 TTNN 支持的内核 FlashAttention，这允许与 TL 进行直接比较。我们专注于非因果变体，它比因果变体暴露出更多的数据流优化机会，因此能更好地说明 TL 的潜力。
- 在评估中，我们将注意力头数量在 64 到 128 之间变化，同时将最终隐藏维度固定为 2048。
- 我们将序列长度从 512 扫描到 8192，并调整批大小，使得总 token 数量固定为 8192。
这确保了所有配置保持在设备的 DRAM 容量内，并使我们能够测量稳态性能而非瞬态效应。
如图 7 所示，在所有测试的头数和序列长度上，TL 始终比 TTNN 带来显著的加速，在几乎所有配置中实现了 1.7 到 2.0 倍的性能提升。这些增益既显著又均匀，表明 TL 的空间映射策略在广泛的注意力形状上具有良好的泛化能力。
其优势源于 TL 自动利用注意力操作数重用结构的能力。TL 放置瓦片，使得关键数据瓦片在多个查询和值瓦片之间在芯片上重用，与 TTNN 的默认映射相比，显著减少了 DRAM 流量，而 TTNN 的默认映射会重复从 DRAM 重新加载这些操作数。性能模型引导 TL 朝着最大化此类芯片上重用、同时遵守建模硬件的片上网络和内存约束的映射方向。
4.3 消融研究
空间重用
为了量化空间重用的效果，我们禁用 TL 的空间重用过程，并强制每个操作数在 GEMM 内核上直接从 DRAM 加载。表 1 报告了两种配置的绝对性能以及 TL 空间重用优化带来的加速比。
随着问题规模增大，效益下降。这一趋势遵循经典的屋顶线模型：对于 GEMM，算术强度随矩阵大小增加，因此较大的问题更可能成为计算受限。一旦内核受限于峰值计算吞吐量，进一步减少 DRAM 流量对端到端运行时的影响就会减弱。
时间重用
图 8 显示了时间重用在各种 GEMM 形状上的影响。时间重用对应于在本地缓冲瓦片并在迭代间重用它们，这样相同的 A 或 B 瓦片不会从 DRAM 重复加载。
我们比较了启用和未启用时间重用的 TL。与空间重用类似，此优化在内存受限区域最有效，因此我们在增加 M 和 N 时减小 K，以保持配置为内存受限。在这些设置下，时间重用产生了最高 1.12 倍的加速比。效益随 M 和 N 增长而增加，因为更大的值增加了这些维度上的迭代次数，从而增加了原本需要重新加载相同 A 或 B 瓦片的次数。
因此，当 M 或 N 很大而 K 很小时，时间重用特别有帮助。对于这种模式不会带来有意义节省的形状，TL 的性能模型会自动降低相应的时间重用映射的优先级，最终选择与没有时间重用的基线相同的映射（和性能）。
性能模型验证
我们通过将 TL 性能模型的预测吞吐量与广泛 GEMM 配置的实测硬件性能进行比较，来验证其准确性。图 9 绘制了模型的估计值和实际测量性能。
平均而言，预测性能与测量值的几何平均差异为 17%。然而，我们的目标不是构建周期精确的预测器，而是一个能够可靠捕捉相对性能趋势的模型，尤其是屋顶线模型所描述的内存受限和计算受限区域之间的转换。如图 9 所示，该模型很好地跟踪了这些转换：它能一致地识别配置何时变为计算受限，并反映了跨形状的性能变化的相对幅度。
静态候选数量 k：如前所述，TL 使用其性能模型对所有候选数据流映射进行排序，选择前 k 个候选，在硬件上剖析这些候选，最后从中选择最佳的一个。参数 k 因此控制着编译成本和最终性能之间的权衡。
为了研究这一权衡，我们将 k 从 1 变化到 5，并在表 2 中报告了跨不同硬件拓扑的几何平均归一化性能和编译时间。这里，“top-1”对应完全静态编译，没有额外剖析（仅使用单个最佳预测映射），而更大的 k 值会向剖析集中添加更多候选。
总体而言，k 的选择对最终性能有适度影响。在 8×8 网格上，top-5 和 top-1 之间的性能差异为 9.3%，其中大部分差距已经从 top-1 到 top-2 时缩小了（提高了 5.8%）。剩余的差异主要源于性能模型偶尔的排名错误，其中预测的最佳映射略差于另一个候选。允许较小的 k（例如 2 或 3）通常足以将真正的最佳映射包含在剖析集中，并恢复接近 top-5 配置的性能，同时保持编译时间适中。
五、相关工作
5.1 硬件建模与协同设计
- 存在关于空间加速器建模与协同设计的工作，例如针对 CGRAs、空间 FPGAs 和脉动阵列，其计算和控制的粒度比 TL 更细。诸如 Spatial、Plasticine、T2S、Halide-to-Hardware 和 HeteroCL 等语言和框架将循环嵌套或功能管道编译到处理单元阵列和本地内存中，通常协同设计覆盖层本身。
- 多面体和张量中心系统，如 AutoSA、TensorLib 和 Rubick，类似地从仿射循环嵌套或张量表达式开始，推导出合成脉动或张量阵列、缓冲区和控制器的时空映射。
- 协同设计工具如 Timeloop/Accelergy 和 AMOS 探索 DNN 加速器的数据流、瓦片化和内存层次结构，而 MAESTRO/MAERI 分析性地建模专用脉动式设计的缓冲区和网络使用情况。
在所有这些系统中，建模单元是内存层次结构中的 PE、缓冲区或循环级别，目标是在该粒度上搜索或合成硬件。相比之下，TL 假设核心内微架构和本地映射是固定的，将每个核心视为原子单元，并仅对核心间的多级内存系统和片上网络进行建模，以决定瓦片实例如何在空间和时间上分布。
5.2 面向空间架构的软件映射与编译
在软件方面，经典的 CGRA 和 FPGA 流程将映射视为 PE 级数据流图上的布局和布线问题：诸如 DSAGEN、面向 ML 的 CGRA 编译器如 ML-CGRA 和 MLIR-toCGRA，以及架构无关的映射器如 Morpher 和 CaSMap，在固定结构上执行布局、布线和模调度。
基于循环和多面体的工具，包括 Timeloop/Accelergy、MAESTRO/MAERI 和 AutoSA，重用循环调度或时空映射作为调度表示，但主要是设计空间探索工具：它们分析性地评估映射，并依赖独立的工具链生成二进制文件。
针对特定空间产品的研究编译器，如 AMOS、LISA 以及针对晶圆级结构的系统级编译器，确实生成每核心程序和通信调度，但通常针对具有固化的映射启发式方法的特定架构系列进行定制。
六、结论
TL 证明，编译器驱动的映射可以在空间数据流加速器上提供厂商级别的性能，同时显著减少手动工作量 。从高级的、以瓦片为中心的内核开始，TL 使用统一捕捉核心、内存和互连的数据流方言，自动选择时空映射、数据重用方案和通信模式。
在 GEMM 和 FlashAttention 等工作负载上，此方法 以最少的人工干预实现了与厂商库竞争甚至更好的性能，将手工设计的内核转变为可重用编译流水线的输出 ，并使更广泛的开发人员能够编写新内核。
同时，TL 的硬件抽象使得编译器过程在很大程度上可在不同加速器间重用 ，简化了对多样化空间数据流架构的支持，并为未来芯片的设计空间探索提供了自然的基底。
关注“鲸栖”小程序，掌握最新AI资讯
本文来自网络搜集，不代表鲸林向海立场，如有侵权，联系删除。转载请注明出处：http://www.itsolotime.com/archives/19356