Dato：突破内存墙，数据流加速器编程新范式实现98%理论峰值利用率

关键词：Dato、数据流加速器、任务型编程模型、通信抽象、虚拟-物理映射

随着人工智能计算的不断发展，数据流加速器将扮演越来越重要的角色。Dato为这些复杂硬件的编程提供了全新的思路，有望成为下一代人工智能计算基础设施的关键组成部分。

Dato的核心创新在于其类型系统的设计，将数据通信和数据分片都提升为一等类型，这使得编译器能够在早期阶段进行深入的分析和优化，而无需依赖复杂且容易出错的手动调度。这种设计哲学不仅提高了性能，还显著降低了编程复杂性，为数据流加速器的广泛采用铺平了道路。

表 1 | 现有加速器编程框架的比较。该表格系统对比了包括CUTLASS、Exo、Triton、Allo、IRON及本文提出的Dato在内的多种主流加速器编程框架。比较维度涵盖目标硬件、输入语言、是否支持显式通信/分片、自动优化能力以及对数据流的类型检查。表格清晰地展示了现有框架在控制力与易用性之间的权衡：低级接口（如CUDA、IRON）控制力强但易用性差；高级语言（如Triton、ARIES）易用但控制力有限。Dato通过引入“流”和“布局”作为一等类型，在保持高级语言易用性的同时，提供了显式的、可静态检查的数据通信与分片控制能力，从而在比较中展现了其独特优势。

论文标题：Dato: A Task-Based Programming Model for Dataflow Accelerators
论文链接：https://arxiv.org/abs/2509.06794

近期的深度学习工作负载，尤其是大语言模型，已将计算需求推至当前内存系统难以承载的程度。许多核心计算（kernel）往往在数据移动、而非实际计算上陷入停滞，形成了众所周知的“内存墙”问题。

尽管现代数据流加速器（如Google TPU、AWS Trainium、AMD Ryzen AI NPU）通过片上流式传输（on-chip streaming）来缓解片外带宽限制，但现有编程模型仍难以有效驾驭这些能力：
* 低级接口提供_细粒度控制却带来高昂开发成本；
* 基于分块（tile-based）的高级语言虽简化编程，却隐藏了通信细节_，限制优化空间并迫使编译器逆向推断数据流。

为此，我们提出了Dato，一种面向数据流加速器的、基于任务（task-based）且嵌入Python的编程模型。

Dato将数据通信与分片（sharding）提升为一等公民的类型结构：开发者通过显式的流类型（stream types）连接任务图编写程序，并使用布局类型（layout types）指定分片输入。任务首先虚拟映射到加速器的空间架构上，随后编译器自动生成符合硬件约束的物理映射。
* 在AMD Ryzen AI NPU与Alveo FPGA上的实验表明，Dato在_显著降低编程负担的同时，实现了高性能；
* 在NPU上，Dato在GEMM运算中实现高达的硬件利用率，注意力核心相比现有商业框架提速；
* 在FPGA上，Dato生成的定制脉动阵列更达到理论峰值性能的，显着领先同类框架。_

本文目录

一、内存墙危机，数据流加速器的机遇与挑战
二、数据流加速器图景，从专用硬件到通用编程挑战
三、Dato编程模型，将数据通信提升为一等公民
- 3.1 设计概述
- 3.2 一个启发性示例
四、流类型，硬件FIFO的直接建模
五、布局类型，自动化分布式计算的利器
六、虚拟映射，从逻辑并行到物理约束的桥梁
- 6.1 虚拟映射图
- 6.2 映射原语
- 6.3 自动化映射流程
七、优化技术，提升硬件效率的关键
- 7.1 内核注入
- 7.2 细粒度布局优化
八、实验结果，从理论到实践的验证
- 8.1 实验设置
- 8.2 NPU上的整体性能
- 8.3 可扩展性
- 8.4 消融研究
- 8.5 可移植性
- 8.6 生产力
结论与未来展望

一、内存墙危机，数据流加速器的机遇与挑战

随着大型语言模型参数规模呈指数级增长，计算模式发生了根本性转变。

过去以计算为中心的架构设计，现在越来越多地被内存受限的工作负载所主导。许多核心计算内核在数据移动、布局转换和同步上停滞，而不是在计算本身。

数据流加速器应运而生，它们通过将计算单元与暂存器和轻量级互连结构耦合，实现了持续的数据流水线处理。即使是通用GPU也采用了类似原则：最近的NVIDIA Hopper和Blackwell GPU引入了专用的张量内存加速器（TMA）来异步流式传输张量。

这些设计通过通道和FIFO将数据移动提升为“一等公民”，能够将内存受限的工作负载转变为高吞吐量的流水线，保持计算单元饱和，同时大幅减少片外通信。

然而，尽管硬件架构取得了这些进展，当前的编程模型仍难以完全释放数据流加速器的潜力，困在性能与生产力的长期权衡之中。

低级别通信控制虽然可以提高性能，但却严重阻碍了生产力。例如，IRON作为AMD NPU的编程前端，向用户暴露了所有FIFO和DMA细节——抽象级别类似于GPU的CUDA。即使是一个简单的矩阵乘法内核也可能需要数百行代码。

基于块的编程语言（如Triton）擅长表达每个块的计算，但对块间或内核间通信的控制很少。像ARIES这样的努力将类Triton语法适配到数据流加速器，但仍然强制执行固定的通信模式，限制开发者进行精细控制。

我们主张，数据通信必须成为数据流加速器编程模型中的一等抽象。

为此，我们提出了Dato——一个基于任务的编程模型，将数据通信和分片都提升为一等类型。在Dato中，程序员定义通过流类型连接的任务，并使用布局类型传入分片数据。这些_任务在空间结构上进行虚拟映射，编译器自动将虚拟映射转换为物理映射_。通过在编程模型级别共同设计计算和通信，Dato明确地暴露任务间流式传输，减少片外通信，并减轻编译器开发负担。

二、数据流加速器图景：从专用硬件到通用编程挑战

要理解 Dato 的价值，首先需要了解现代数据流加速器的架构特点及其编程挑战。

脉动阵列架构已成为现代深度学习加速的基石，能够为密集线性代数工作负载维持高吞吐量。

专用脉动阵列：如 Google TPU、AWS Trainium 和 Inferentia 等架构，使用大型处理元素网格来流式传输矩阵乘法，在本地重用操作数，并将部分结果传递给邻居，从而大幅减少片外内存通信。
可编程数据流结构：将脉动原理扩展到更多样化的工作负载，通过对片上通信和计算提供更精细的控制。例如，Cerebras WSE 通过定制互连结构连接整个晶圆上的数十万个核心，实现大规模并行性。
异构 AI 引擎：如 AMD 的 Ryzen AI NPU，采用了 VLIW 处理器和向量单元组成的 AI 引擎阵列，通过可编程片上网络连接，支持流式和通用任务。

图 1 | 典型的数据流加速器架构 [46,54]。硬件单元的实际数量可能更多。该架构包含计算单元阵列、本地内存、流接口、片上交换盒、暂存器内存及 DMA 引擎，可通过模块间直接流传输避免中间数据绕经片外 DRAM，减少内存墙影响，适配深度学习等内存密集型工作负载，提升数据流动与计算效率。

图 1 展示了一个典型的数据流架构：一个包含本地内存和流式接口的计算单元阵列，与邻居交换数据。片上交换箱在整个阵列中路由流量，而暂存器内存用于缓冲张量并支持跨块的并行访问。

表 1 | 现有加速器编程框架的比较。该表格系统对比了包括 CUTLASS、Exo、Triton、Allo、IRON 及本文提出的 Dato 在内的多种主流加速器编程框架。比较维度涵盖目标硬件、输入语言、是否支持显式通信/分片、自动优化能力以及对数据流的类型检查。表格清晰地展示了现有框架在控制力与易用性之间的权衡：低级接口（如 CUDA、IRON）控制力强但易用性差；高级语言（如 Triton、ARIES）易用但控制力有限。Dato 通过引入“流”和“布局”作为一等类型，在保持高级语言易用性的同时，提供了显式的、可静态检查的数据通信与分片控制能力，从而在比较中展现了其独特优势。

从表 1 可以看出，现有编程框架在数据流类型检查和显式分片支持方面存在明显不足。这正是 Dato 试图解决的核心问题。

三、Dato 编程模型：将数据通信提升为一等公民

Dato 的设计哲学是让数据通信成为编程模型中的一等抽象。为实现这一目标，Dato 采用了基于任务的编程范式，将计算组织为通过显式流类型连接的任务图。

3.1 设计概述

Dato 被设计为一种嵌入在 Python 中的加速器编程语言，旨在最大化可用性和生产力。它建立在开源加速器编程框架 Allo 之上，并重用了 MLIR 基础设施，同时引入了针对数据流加速器编程定制的显著新功能。

图 2(a) | Dato 编译器概述。编译器流程始于用户编写的 Python 函数，先经数据流特定类型检查生成带类型注释的抽象语法树（AST），再通过 MLIR 构建中间表示（IR），后续经虚拟-物理映射、内核注入等优化，最终生成适配 NPU 或 FPGA 的目标代码，兼顾性能与可移植性。

如图 2a 所示，Dato 编译器流程始于用户编写的 Python 函数，经过数据流特定类型检查生成带类型注释的抽象语法树（AST），再通过 MLIR 构建中间表示（IR）。随后，编译器生成捕获程序数据流结构的虚拟计算图。

由于数据流加速器通常具有有限数量的处理引擎，而用户程序可能表达更大的计算，Dato 会执行虚拟到物理的映射步骤。这一步骤弥合了逻辑并行性和有限硬件资源之间的差距，实现在空间结构上的高效任务放置和调度。

物理映射之后，编译器应用内核注入和布局优化通道来优化内存放置和计算操作。最终优化的 MLIR 模块随后被转换为特定于后端的代码：
* 针对 AMD NPU，利用 MLIR-AIE 作为后端；
* 对于 FPGA，则为高级综合生成 C++代码，并使用 AMD Vitis 进行综合。

3.2 一个启发性示例

我们通过一个简单的生产者-消费者模式来说明 Dato 的优势。图 2b 展示了用 Allo 编写的代码，而图 2c 展示了相应的 Dato 实现。

图 2(b) | 用于对比的 Allo 程序示例。该示例需采用数组式编程，手动通过 .relay() 完成数组-流转换，且需为每个处理单元（PE）编写单独函数如 producer_0、producer_1，存在语义间接性、索引管理繁琐等问题，凸显 Allo 在数据流编程中的局限性。

观察图 2(b) 的 Allo 程序示例，可以发现几个限制生产力和可扩展性的问题：
1. 语义间接性：在 Allo 中，即使用底层硬件通过 FIFO 通信，用户也必须以基于数组的风格编写程序，并需使用 .relay() 原语显式插入数组到流的转换。这造成了语义鸿沟：程序员基于数组进行推理，而硬件执行流。
2. 繁琐的索引管理：Allo 代码需要显式分区和索引切片。考虑到不同的张量可能对不同的计算单元有不同的分片方法，这种低级处理容易出错，在大型程序中难以管理。
3. 缺乏可扩展性：Allo 要求为每个 PE 编写单独的函数（例如 producer_0、producer_1）。虽然对于小示例可行，但这种方法难以扩展到具有数百个 PE 的加速器。

图 2(c) | Dato 程序示例。示例中 Dato 将流作为一等类型，通过 @task 定义任务并指定虚拟映射，利用布局类型自动处理数据分片，无需手动编写多 PE 函数，解决了 Allo 的语义间接性与重复编码问题，大幅简化数据流编程并提升代码可维护性。

为解决上述限制，Dato 引入了基于任务的编程模型，如图 2c 所示。用户编写一个顶层函数，在其中可以使用 @task 定义多个并行任务。
* 流作为一等类型：Dato 将流提升为一等类型，允许程序显式表达硬件级通信，解决了限制一（语义间接性）。
* 自动布局与分片：分片输入通过布局细化类型指定，使编译器能够自动处理分区和布局转换，有效消除了手动索引运算，解决了限制二（繁琐的索引管理）。
* 虚拟映射：用户定义单个任务并应用虚拟映射，让编译器在硬件上复制和调度它，从而克服了限制三（缺乏可扩展性）。

通过使抽象与硬件语义对齐并自动进行布局和映射，Dato 消除了 Allo 所需的间接性和簿记，简化了数据流加速器的编程。

四、流类型：硬件 FIFO 的直接建模

我们在前端暴露了一等流类型，以直接建模片上通信，消除数组到流的转换。流处理中的先前工作通常将流抽象为无界的高级语义结构，或在多生产者和多消费者设置中，与硬件容量约束解耦。现有的面向硬件的框架尝试形式化流，但通过功能或时空类型实现，依赖于运行时握手或调度对齐，而没有静态保证。

相比之下，我们的流类型嵌入在命令式 Python 前端中，并忠实建模硬件 FIFO 的语义，且可推广到标量 FIFO 之外。

流本质上是一个先进先出的结构，它定义了两个任务之间的点对点连接。其类型定义为 Stream[T, N, P]，其中 T 是元素类型，N 指定以元素为单位的逻辑容量，P 定义每次传输捆绑的元素数量（默认情况下不进行打包）。.put() 和 .get() 是与该类型关联的两个操作，用于实现任务间的通信，如图 2c 的 L16 和 L21 所示。

图 2(c) | Dato 程序示例。示例中 Dato 将流作为一等类型，通过 @task 定义任务并指定虚拟映射，利用布局类型自动处理数据分片，无需手动编写多 PE 函数，解决了 Allo 的语义间接性与重复编码问题，大幅简化数据流编程并提升代码可维护性。

图4 | 流类型的类型规则。Γ是普通的类型上下文，而Δ是一个多重集，表示跟踪必须被恰好使用一次的资源的线性上下文。⊎是多重集并集。𝑁是流的深度。

图 4 展示了流的类型规则，它精确地反映了硬件 FIFO 的语义。每个元素占据头（出队）和尾（入队）之间的一个槽位，我们使用线性能力令牌 Free(S) 和 Used(S) 来跟踪这些槽位：
* 仅当 Free(S) 令牌可用时，.put() 操作才是良类型的，该操作会将 Free(S) 转换为 Used(S)（对应尾指针前进）。
* 对称地，.get() 操作仅当存在 Used(S) 令牌时才是良类型的，它会立即将 Used(S) 转回 Free(S)，同时返回一个线性就绪未来作为有效负载。这是一个无需阻塞即可用但必须被恰好使用一次的值。

.get() 操作在消费者接受项目时即释放槽位，这缓解了反压，使得生产者即使在消费者延迟使用的情况下也能继续流式传输。Dato 有意将数据移动与实际计算分离，并引入内部的 .await() 操作来解包未来，即获取要使用的实际值。由于未来是就绪且线性的，编译器可以将唯一的 .await() 放置在延迟隐藏点，而不会引发重复或死锁的风险。

循环令牌强制执行一个简单的不变量：每个 .put() 将 Free(S) 减一并将 Used(S) 加一，每个 .get() 则相反。因此，溢出和下溢在构造上变得不可类型化。.await() 操作是 FIFO 中性的，它仅解包一个已经出队的元素，同时保持了安全性和吞吐量。

为了对使用流的程序进行类型检查，Dato 使用 N 个 Free(S) 令牌和零个 Used(S) 令牌的多重集来初始化每个 Stream[T, N]。在每个程序点，系统维护一个将流映射到其当前令牌多重集的抽象状态。类型检查通过控制流图上的前向抽象解释来执行。如果在退出每个区域时所有未来都被消费且没有令牌泄漏，则程序是良类型的。

图 5 | 未通过类型检查的数据流程序示例。示例包含两类错误：一是 .put() 与 .get() 操作不一致（如循环内 put 次数超流容量），二是死锁（如两任务相互等待对方流数据）。Dato 的类型系统通过跟踪流的 Free/Used 令牌，在编译时拦截这类错误，保障数据流程序的正确性与安全性。

图 5 展示了几个被类型系统拒绝的不可类型化程序。需要注意的是，该类型系统是一个无定时模型，因此无法提供依赖于具体硬件调度的定时信息，例如最小 FIFO 深度。

通过使流成为 Dato 中的一等类型，用户可以自然地表达更高级的模式，如流数组或张量流。这种可组合性可以扩展到多级流水线和分叉-连接拓扑，实现清晰且可分析的通信，同时保留到硬件的直接映射。编译时检查不仅保证了安全性，还提供了早期反馈，有助于性能工程师调试和优化其设计。

五、布局类型：自动化分布式计算的利器

随着深度学习扩展到分布式异构系统，将单个处理单元或设备作为一等编程单元来处理很快变得难以管理。因此，单程序多数据范式成为普遍基础，它表达在分区数据上运行的一个程序，同时由系统协调放置和通信。

基于 SPMD 范式，近期关于注释分片的工作提供了（半）自动化机制，用于为大型模型分区张量和调度计算。类似的挑战也出现在拥有数百个处理单元的数据流加速器上。如第 3.2 节所讨论的，手动计算每个数据块的索引运算并确保内核级一致性是脆弱且耗时的。

我们通过将数据布局提升为一等静态构造迈出了下一步：跨处理单元的布局直接编码到类型系统中，作为基础张量类型的细化。在 Python 中，这表示为 T[Shape] @ Layout("...")，如图 2c 中的 L14 和 L19 所示。

图6 | 布局类型定义和最小类型规则。𝐷是基本元素域。Γ是类型上下文。𝜏!Π表示类型𝜏的待处理集合效应Π。

图 6 定义了布局类型及其类型规则。每个张量轴都携带一个布局标签，要么是 R，要么是 S。

图 3. 【(a) 图 2c 中程序的映射情况】。该映射展示 Dato 程序中生产者任务与消费者任务的对应关系：生产者将输入数据 A 分片后通过流 Z 传递给消费者，消费者处理后输出到 B。此过程体现 Dato 通过虚拟映射将逻辑任务高效映射到硬件 PE 的能力，减少数据移动开销。【(b) 16×16 张量映射到 2×2 网格的布局】。布局中张量维度标注 “S（分片）” 或 “R（复制）”，16×16 张量按 2×2 网格分片后，每个网格单元处理对应子张量。Dato 通过布局类型静态检查分片合理性，确保数据在硬件上的高效分配与访问，适配空间并行计算。

图 3b 展示了一个二维张量的布局示例。沿轴 i 的约简操作在两种情况下都会从形状中消除该轴，但单个分片上可用的信息不同：
* 如果 L[i] = R，每个分片已经保存了轴 i 的完整切片；因此约简纯粹是局部的和最终的，所以效果集保持不变。
* 如果 L[i] = S，每个分片仅看到轴 i 的分区；局部约简产生部分结果，仍然必须跨分片组合。我们通过将运算符 ⊕（例如 +、max）添加到效果集 Π 中来记录这一义务。

待处理的效果由不改变形状或布局的集体操作消除。allreduce⊕ 的规则简单地从 Π 中移除 ⊕，意味着每个分片的部分结果已全局组合，并且该值现在相对于 ⊕ 具体化。这种设计清晰地分离了算法意图（轴上的张量约简）和分布同步（跨分片组合），并保持了可组合性：部分结果可以流经后续操作，累积或保存在 Π 中，直到消费者需要具体化的值。

默认情况下，布局类型是 R，表示所有维度都是复制的。这意味着用户可以编写普通的张量代码，除非需要，否则无需额外的布局注释。将布局附加为细化类型实现了布局兼容性的静态检查，防止不安全或隐式的重新分片，并允许 Dato 在出现不一致时自动插入集体操作并提供精确诊断。这种设计保留了张量程序的熟悉语法，同时用布局丰富了类型系统，以最小的开销为用户提供快速且可操作的反馈。

六、虚拟映射：从逻辑并行到物理约束的桥梁

在本节中，我们介绍用于扩展设计的映射原语，并讨论自动化虚拟到物理映射过程的技术。

6.1 虚拟映射图

当单个分块维度超过片上容量时，Dato 自然地推广到多维分块和映射。

图 7(a) | Dato 中的分块 GEMM 内核。该内核通过 Layout 类指定张量 A、B、C 的分片方式（如 LyA 为 “S1S2”），用 @task 定义 GEMM 任务并设置 3D 虚拟映射，内部先计算部分结果 part_C，再通过 allreduce 聚合得到最终结果，适配大规模并行计算，提升 GEMM 任务性能。

如图 7a 的 L5-7 所示，布局注释指定了每个张量维度如何分片以及映射到哪个设备轴，指示该维度在空间结构中的最终放置。例如，S152 表示张量 A 的维度 0 被分片并映射到设备轴 1，而维度 1 被分片并映射到设备轴 2。

给定映射指令（L10），Dato 将任务实例化为一个 3D 虚拟网格，网格中的每个点对应一个任务实例。生成的虚拟映射如图 7b 所示。

图 7(b) | GEMM 内核的虚拟映射扩展：虚拟映射将 GEMM 任务实例化为 3D 网格结构，每个实例处理部分张量数据。任务间通过流传递中间结果，集体操作（如 allreduce）对应特定的映射关系。此映射桥接了逻辑并行性与硬件资源，为后续物理映射奠定基础。

为了捕获实际计算并将虚拟计算折叠到物理 PE 上，我们根据程序的流操作和布局注释连接任务实例，构造虚拟映射图（VMG）。VMG 中的每个节点代表一个任务体，即一个 PE 级内核（例如，局部矩阵乘法）。有向边由通信原语引入：Stream 上的每个 .put() 或 .get() 操作会建立一条从生产者到消费者的边。

图 7(c) | 分块 GEMM 内核虚拟到物理映射中的 .bundle() 和 .chain() 原语。.bundle() 将同构节点（如 A0B0 与 A1B1 乘法节点）合并为多轮执行节点，减少物理节点数量；.chain() 融合生产者-消费者节点（如乘法与加法节点）为顺序节点，消除中间流缓冲。二者在保持语义的同时优化硬件资源利用，适配不同加速器约束。

集体运算符如 .allreduce() 会生成约简边，连接所有参与该集体操作的任务实例，具体连接方式由布局指示。例如，当约简维度沿虚拟轴分片时，.allreduce() 会将部分结果（如 part_C）聚合成最终输出（如 C），如图 7a 的 L14-16 所示。生成的 VMG 是与设备无关的，其结构在图 7c(i) 中可视化，描绘了以特定节点为根的子树。以其他节点为根的类似子树共享相同的结构。

6.2 映射原语

为了将 VMG 中的每个节点映射到硬件上的实际物理 PE，我们引入了两个调度原语：.bundle() 和 .chain()。它们在保持程序语义的同时重塑图拓扑。这些原语可向用户公开以进行手动控制，也可由 Dato 自动应用。

.bundle() 原语：将一组同构节点（即具有相同计算和 I/O 模式的节点）合并为单个多镜头节点。合并后的节点被调度执行多次，依次处理每个原始任务的输入或输出。从概念上讲，.bundle() 将显式的扇入/扇出转换为时分复用的连接或分割。这减少了物理节点和外部端口的数量，消除了中间缓冲区。在路由或端口可用性有限但计算能力足以在同一处理块上托管多次调用的硬件上，此方法尤其有效。
.chain() 原语：将两个相连的节点（通常是生产者-消费者对）融合为单个顺序节点。融合后的节点保留所有外部边，并按数据依赖性排序计算。这通常移除了中间流及其关联的缓冲。通过将路径折叠到单个位置，.chain() 提高了数据局部性并减少了路由压力。融合后的节点仍需遵守平台的端口约束（例如，最多两个全局输入和输出，或兼容类型的共享端口）。

图 7(c) | 分块 GEMM 内核虚拟到物理映射中的 .bundle() 和 .chain() 原语。.bundle() 将同构节点（如 A0B0 与 A1B1 乘法节点）合并为多轮执行节点，减少物理节点数；.chain() 融合生产者-消费者节点（如乘法与加法节点）为顺序节点，消除中间流缓冲。二者在保持语义的同时优化硬件资源利用，适配不同加速器约束。

图 7c 展示了一个部分和模式的示例：
* 从 (i) 中的初始 VMG 开始，对并行乘法节点应用 .bundle() 将它们的输出合并为单个连接节点，将结果累积到中间变量中。
* 如 (ii) 所示，进一步与加法器应用 .chain() 产生单个融合节点。
* 该融合节点顺序计算并在内部执行累积，如 (iii) 所示。

在外部，此图仅暴露两个全局输入（在镜头间重用的 A 和 B 流）和一个输出。此变体以减少节点间并行性为代价，最小化了路由和缓冲开销。

6.3 自动化映射流程

当 VMG 很大时，手动将图折叠到物理 PE 上变得繁琐且容易出错。为此，我们提出了一种搜索算法，用于生成映射原语序列及相应的物理映射程序。

算法1 | Dato编译器自动化虚拟到物理映射的核心搜索算法。它以一个代表程序逻辑并行度的虚拟映射图和硬件物理处理单元数量为输入，通过系统地应用 .bundle()（捆绑）和 .chain()（链式融合）两种映射原语，递归地探索不同的节点合并策略。搜索过程受限于目标平台的资源约束（如I/O端口数量、内存访问限制）和依赖关系合法性检查，目标是找到一种在满足硬件限制的同时，能最大化性能的物理任务布局方案，从而将用户声明的高级任务图高效地部署到有限的空间计算架构上。

如算法 1 所示，该算法通过系统地“并置”节点，将 VMG G 映射到具有资源预算 C 的计算结构上，直到虚拟节点的数量落在可用预算内。

映射原语仅在满足轻量级、架构感知的合法性条件时才被应用。这些条件主要受三个标准约束：
1. 资源约束：例如顶级 I/O 端口可用性、NPU 上每块的入口和出口限制，以及 FPGA 上的 BRAM 访问限制。
2. 接口兼容性：要求节点具有匹配的元素类型、形状和生产者-消费者关系，以确保外部行为等价。
3. 依赖性健全性：仅允许那些独立或直接由生产者-消费者边连接的节点被合并，以避免由跨传递路径折叠可能引起的循环。

对于 I/O 核算，只有全局端口受到约束。并置节点内的本地通信被视为内部通信，不计入全局端口预算（例如，NPU 上的块本地链接或 FPGA 中的片上布线）。在这些规则下，每个并置步骤在遵守端口、路由和内存的平台约束的同时，减少了同时活动节点的数量，使算法在不同目标之间具有可移植性。

在高层，该过程在以初始 VMG 为根的树上执行分支-约简搜索。树中的每个节点代表一个由原语应用序列产生的映射图（或状态）。例如，图 7c 中所示的三个图说明了此搜索树内的不同状态。每条边对应一个合法的原语应用，它严格减少虚拟节点的数量，确保单调下降。从任何中间状态，算法枚举候选节点对，并使用合法性条件进行过滤。允许的对产生子状态，这些状态被递归探索。

虚拟节点数不超过 C 的状态被记录为可行映射，而那些仍超出资源预算且不允许进一步有效转换的状态被剪枝。为了指导探索，我们引入了一个优先级函数，该函数偏爱可能提高硬件兼容性的转换（例如，减少全局 I/O 使用或缓解端口压力）。这种启发式排序提高了收敛速度而不牺牲正确性。一旦找到可行映射，候选者将经过特定于目标的优化和代码生成。对映射进行评估，直到获得所需数量的有效实现。最终，选择具有最高测量性能的映射作为输出。

七、优化技术：提升硬件效率的关键

本节讨论数据流加速器必要的进一步优化。

7.1 内核注入

许多 NPU 为其 VLIW 核心提供手动调优的 C++ 向量内核，而 FPGA 提供预验证的 C++ HLS IP 核以供重用。为了利用这些库函数，我们定义了一个具有显式契约的内核接口：
1. 可允许的形状和向量宽度。
2. 所需的 I/O 布局（例如，分块或打包）。
3. 成本/资源提示，如延迟和启动间隔——针对每个目标增加设备规格。

当计算图满足契约时，Dato 的降低通道会选择适当的变体并合成一个包装器，将任务循环分块到内核的微分块，插入前奏/结尾布局适配器，并拼接流和 DMA 以重叠打包/解包与计算。

7.2 细粒度布局优化

我们进一步引入了细粒度布局优化，它在单个处理单元（PE）内运行，因为高性能内核通常对输入/输出布局有严格的要求。

我们遵循 CuTe 和 Graphene 的思路，通过记录元组（偏移量、大小、步长）来显式表示张量布局。在注入的内核周围进行简单的降级操作，往往会留下一系列小的重新格式化步骤链——例如分块、打包、转置、重新打包——这增加了带宽开销并给本地内存带来压力。

为解决此问题，我们开发了两种互补的优化技术：规范化与折叠，以及 DMA 感知的提升。

规范化与折叠：我们将布局变换建模为布局代数中的元素：即受分块可除性约束的整数域上仿射映射的组合。Dato 将每个 MLIR transform_layout 操作序列提升到此代数中，对组合进行规范化（例如，合并相邻块，取消逆对），然后将其折叠为最小范式。具体来说，我们在静态单赋值（SSA）形式中维护布局元组，符号化地组合它们，并通过轻量级算术检查来解除完整性侧条件。结果消除了冗余的具体化，并将兼容的步骤融合为单个打包/解包循环或单个 DMA 描述符。与局部窥孔优化相比，代数规范化能够暴露跨循环边界以及生产者-消费者任务之间的全局取消机会。
DMA 感知的提升：许多布局操作，如打包、交错、步长和固定模式排列，可以由 NPU 的 DMA 引擎在数据传输过程中（in-flight）执行。我们的通道在流边界提升符合条件的布局工作，并将其折叠到 DMA 描述符中：例如，打包变为跨步收集，转置变为 2D 步长交换。提升过程受到对步长范围、突发大小和存储体映射的合法性检查的保护，从而提高了有效带宽并释放了本地 SRAM。在 FPGA 后端，相同的变换被降级为 HLS 友好的突发对齐循环。最后，调度器将这些 DMA 操作放置在生产者和消费者发起的时隙中，并依赖多缓冲机制在引擎执行变换时保持计算单元的饱和。

八、实验结果：从理论到实践的验证

在本节中，我们首先讨论实验设置，并评估我们的编译流程。

8.1 实验设置

我们在 NPU 和 FPGA 上评估 Dato，它们分别代表了商业化的粗粒度和细粒度数据流加速器。NPU 具有针对超长指令字（VLIW）执行优化的固定功能 PE 阵列，而 FPGA 提供基于查找表（LUT）的可重配置性，支持自定义数据流结构。

对于 NPU 评估，我们与 IRON（基于 MLIR-AIE 的官方 AMD 流程）和 ARIES（最先进的 AIE 编译框架）进行比较。实验在具有 20 个计算块（4×5 阵列）的 AMD Ryzen AI NPU（XDNA1）上进行。

对于单内核基准测试，我们报告设备内核执行时间，不包括一次性主机到设备的数据传输延迟，以更好地隔离硬件利用率——硬件利用率定义为实际计算吞吐量与理论峰值的比率。
对于多内核工作负载，我们测量端到端延迟，包括数据移动开销。

每个基准测试包括 20 次热身迭代，然后是 400 次定时运行，报告算术平均值。对于 FPGA 评估，我们以 AMD Alveo U280 板卡为目标，并与 Allo 框架进行比较。设计使用 Vitis 2023.2 进行综合，延迟数据取自 HLS 报告。

8.2 NPU 上的整体性能

我们首先测量 Dato 在单内核和多内核设计上的整体性能。

8.2.1 单内核设计

单内核通用矩阵乘法（GEMM）基准测试扫描矩阵大小 M, N, K ∈ {256, 512, 1024, 2048}（共 64 种配置）。我们将内核映射到 NPU 上的 4×4 活动块区域，并评估三种精度：int16（i16）、bfloat16（bf16）和 int8（i8）。每种设置的分布是固定的：对于 i16/bf16，我们使用 m = n = k = 64；对于 i8，我们使用 m = k = 64 和 n = 128。

图 8 | 不同数据类型下 GEMM 内核的性能结果。阴影区域表示执行方差。结果显示，Dato 在 i8 精度下硬件利用率达 61.58%，超 IRON 和 ARIES；i4+i8 混合精度利用率最高达 84.38%；bf16 和 i16 精度也接近或超越对比框架。且执行方差小，证明 Dato 在不同精度下适配 GEMM 任务的高效性与稳定性

图 8 展示了屋顶线模型。在所有精度下，吞吐量随着算术强度的增加而增加，然后在达到计算上限时饱和。

在 i16 精度上，Dato 匹配了最佳供应商流程的性能，达到了峰值性能的 75.01%。
在 bf16 精度上，Dato 达到了 47.56%，略低于 IRON 但高于 ARIES。

最大的性能差距出现在 i8 精度上，Dato 达到了 61.58% 的利用率，优于 IRON 和 ARIES；后者由于缺乏 DMA 调度，在高强度下出现了明显的性能差距。这些趋势在整个算术强度范围内保持一致，Dato 保持或接近从带宽受限区域到平台期的屋顶线顶部曲线。

结果表明，Dato 的优化可以匹配甚至超过 IRON 中手动编写模板的性能。我们将这一优势归因于三个因素：

（i）虚拟映射，能够随着分块大小干净地扩展并平衡数千个任务，这是大型 GEMM 的先决条件；
（ii）通信感知的 DMA 调度与活跃性指导的端口重用，在 NPU 的 2 输入/2 输出端口约束下维持了数据供应；
（iii）向量化微内核与布局规范化/折叠和 DMA 感知提升的结合，使循环分块与内核的向量宽度对齐，并最小化了打包/转置开销。

此外，Dato 启用了其他框架不支持的混合精度功能。如图 8 所示，使用 i4 和 i8 张量允许更大的分块配置，最终通过充分利用较低位宽实现了高达 84% 的硬件利用率。

8.2.2 多内核设计

我们评估了两个代表 Transformer 工作负载的多内核设计。第一个是多头注意力（MHA），实现为两个 GEMM 之间夹着一个数值稳定的 softmax；第二个是前馈网络（FFN），其中我们实现了 LLaMA 风格 SwiGLU 块中使用的上和门投影。我们还在 Dato 中实现了 FlashAttention 内核，据我们所知，这是首次在 AMD NPU 上使用高级编程框架演示 FlashAttention。

我们设置序列长度 L ∈ [128, 2048]，具有 12 个头和头维度 64；FFN 使用隐藏大小 768 和 3072。由于 ARIES 不提供可用的多内核实现，因此我们仅将 Dato 与 IRON 进行比较。所有实验在 NPU 上使用 bf16 精度，激活 4×4 计算块。由于 IRON 暴露的是固定的、每内核模板，内核必须顺序调用，这强制了阶段之间的临时数据写回。

图 9 | 多头注意力（MHA）和前馈网络（FFN）在 NPU 上的性能（bf16 精度）。Dato 实现的 MHA（含 FlashAttention）在序列长度 2048 时较 IRON 提速 2.81 倍，FFN 提速 1.64 倍。因 Dato 支持片上流传输与任务融合，消除核间数据写回 / 读取开销，而 IRON 需顺序调用内核，凸显 Dato 在多内核 Transformer 任务中的优势

多头注意力

为了充分利用数据流架构的优势，我们在 NPU 上实现了 FlashAttention 内核，这是一种使用在线 softmax 的融合 MHA，实现了流式 MHA 而无需将中间结果写回 DRAM。

图 10 | 左：Dato 中 FlashAttention 内核的实现。右：注意力内核的物理映射。左侧实现使用流类型连接 send_q、gemm0 等任务，并通过布局类型管理张量分片；右侧物理映射将任务分配到 NPU 处理单元，实现数据流式处理，无需将中间数据写回片外内存。这是 Dato 首次在 AMD NPU 上通过高层编程框架实现 FlashAttention 并取得高性能的关键。

图 10 展示了 Dato 的实现及其物理映射，其中流类型和布局类型使程序更易于编写和阅读。

图 9 | 多头注意力（MHA）和前馈网络（FFN）在 NPU 上的性能（bf16 精度）。

如图 9 所示，Dato 在不同序列长度上始终优于 IRON，最大加速比达到 2.81 倍。性能增益主要源于：
（i）通过片上流式传输和空间扩展减少了片外通信，消除了内核之间的数据写回与读回；
（ii）通过跨块运行单个融合流水线，降低了内核启动开销；
（iii）更小的每块工作集，使 Dato 能够处理超过 2048 的序列长度，而 IRON 因 DMA 传输限制而无法支持。

前馈网络

在 Dato 中，我们将两个 GEMM 操作空间融合到单个数据流流水线中，将 4×4 处理阵列划分为两个 4×2 子网格，分别用于一个 GEMM。相比之下，IRON 无法融合这些内核：它顺序调用两个独立的 GEMM 内核，强制要求阶段之间的全局内存写回/读回，并产生额外的每内核启动开销。

图 9 | 多头注意力（MHA）和前馈网络（FFN）在 NPU 上的性能（bf16 精度）。Dato 实现的 MHA（含 FlashAttention）在序列长度 2048 时较 IRON 提速 2.81 倍，FFN 提速 1.64 倍。Dato 的优势在于支持片上流传输与任务融合，消除了核间数据写回/读取开销，而 IRON 需顺序调用内核，这凸显了 Dato 在多内核 Transformer 任务中的优势。

如图 9 所示，这种融合执行相对于 IRON 带来了 1.64 倍的加速。增益主要源于减少的片外通信、改进的流水线占用率以及更高的块利用率。

8.3 可扩展性

图 11 报告了随着活动网格从 1×4 增加到 2×4 再到完整的 4×4 阵列，NPU 上 GEMM 的可扩展性。

图 11. 左：GEMM 的可扩展性结果；右：GEMM 的优化消融实验，柱形上的数字表示 NPU 执行时间。

左图：当 NPU 活跃网格从 1×4 增至 4×4 时，GEMM 延迟几乎按比例降低，4×4 配置较 1×4 提速 3.67 倍。尽管由于通信和规约开销增加，硬件利用率略有下降，但结果证明 Dato 的虚拟映射、布局优化等机制能有效支持任务随硬件资源扩展，适配大规模计算需求。

右图：基线 NPU 时间为 13.65ms，启用虚拟映射后降至 6.42ms（降幅最大），注入向量化内核后达到 3.13ms，再经布局优化后最终为 2.83ms。数据表明，虚拟映射有效减少了编排开销，而其他优化则提升了硬件利用率，共同构成了 Dato 高性能的核心优化栈。

对于最大的工作负载，绝对延迟几乎与块数量成比例地减少：2×4 配置相对于 1×4 实现了 1.97 倍的加速，而 4×4 配置达到了 3.67 倍。尽管由于通信和规约开销增加，利用率略有下降，但结果表明 Dato 能够随计算块数量有效扩展，突出了其虚拟映射、布局优化和 DMA 调度在保持大规模利用率方面的重要性。

8.4 消融研究

图 11 展示了 Dato 优化堆栈的消融实验，将延迟分解为 NPU 计算、数据/指令复制和环境设置。

从基线开始，NPU 时间为 13.65 毫秒；
启用虚拟映射将其减少到 6.42 毫秒，同时折叠了主导的设置开销，产生了最大的端到端增益。
注入向量化内核将计算时间减半至 3.13 毫秒（-51%）。
随后的布局优化将其进一步修剪至 2.83 毫秒（-10%）。

在整个优化序列中，数据/指令复制成本始终可以忽略不计，总延迟相对于基线下降了约 4 倍。

这些结果表明，虚拟映射主要消除了编排和启动成本，而其他优化则通过增加块利用率和减少变换开销来驱动剩余的改进。

8.5 可移植性

为展示 Dato 在 NPU 之外的可移植性，我们生成了一个 16×16 i8 输出驻留脉动阵列，运行在 300 MHz，以 AMD Alveo U280 FPGA 为目标，这是一个截然不同的硬件平台。用户可以保留相同的 Dato 前端，并简单地指定 FPGA 作为构建目标。

图 12. 左：Dato 生成的脉动阵列架构。右：Alveo U280 FPGA 上的性能结果。标注数字为实现的硬件利用率。

左图中的架构包含计算单元（COMP）、多级缓存与流接口，各处理单元通过 Dato 流互联成菊花链结构，数据按脉动方式流动，实现了计算与数据传输的重叠。该架构由 Dato 任务封装组件并自动生成，适配 Alveo U280 FPGA 的硬件特性。

右图显示，Dato 在 GEMM 任务上达到 150 GOP/s，接近理论峰值 153.6 GOP/s，且设计满足 300MHz 时序目标；而 Allo 因扇出过大仅达到 132MHz。Dato 性能较 Allo 最高提升 8.2 倍，证明了其在 FPGA 上的高性能与硬件适配性。

我们实现了一个菊花链脉动阵列架构，如图 12 所示，其中每个 PE 和多级缓存被封装为 dato.task 模块，并使用 Dato 的流进行互连。该后端的完整实现如下。

一种菊花链 systolic 阵列。下面展示了菊花链 systolic 阵列的 Dato 实现，其中 dato.meta_if 和 dato.meta_else 是编译时条件分支，不会产生硬件开销。

对于此后端，编译器绕过了 NPU 特定的 DMA 优化，而是生成用于 FPGA 综合的 C++ HLS 代码。

图 12. 左：Dato 生成的脉动阵列架构。右：Alveo U280 FPGA 上的性能结果。标记的数字是实现硬件利用率。

我们评估了从到的 GEMM 工作负载，并与最先进的框架 Allo 进行比较。图 12 显示 Dato 始终优于 Allo，最高可达 8.2 倍。在最大的矩阵尺寸下，Dato 实现了 150 GOP/s，接近脉动阵列的理论峰值 153.6 GOP/s（假设每个 PE 每个周期执行一次 MAC 操作）。

实现这种性能的关键在于 Dato 能够通过精心设计的多级内存层次结构来重叠数据传输和计算，这是 Allo 由于其有限的原始语集合而无法从简单 GEMM 内核实现的转换。

8.6 生产力

最后，我们比较了代码行数方面的实现工作量。

表 2 | 不同内核的代码行数比较。该表格对比了在实现相同功能（GEMM、MHA、FFN内核）时，底层框架 IRON 与高级框架 Dato 所需的代码行数。数据显示，Dato 仅需极少的代码（分别减少至 IRON 的 1/12、1/5.4、1/7.2）即可达到与 IRON 相当甚至更好的性能。这直接量化了 Dato 在提升程序员生产力方面的巨大优势。其根本原因在于 Dato 的编程模型：用户无需手动管理数据分片索引、FIFO 通信细节或为每个处理单元重复编写代码，而是通过声明式的任务、流和布局类型，由编译器自动处理这些繁琐且易错的底层细节，从而大幅简化了优化代码的编写过程。

如表 2 所示，Dato 在匹配 IRON 性能的同时，代码行数减少了 12 倍（参见图 7a 和 10 中的示例）。除了简洁性之外，Dato 还解锁了空间多内核融合和无缝重定向到其他加速器的能力——这些能力显著提高了性能工程师的生产力。

结论与未来展望

在本文中，我们介绍了 Dato，这是一个 Python 嵌入的、基于任务的编程模型，它将数据通信和分片提升为一等类型，并将虚拟映射编译为数据流加速器上的高效物理部署。Dato 通过其流类型和布局类型的创新设计，实现了硬件抽象与编程模型之间的紧密对齐，为数据流加速器编程提供了一种既高效又易于使用的解决方案。

我们的实验结果表明：

Dato 在 AMD Ryzen AI NPU 和 Alveo FPGA 上都实现了卓越的性能，同时显著降低了编程复杂性。
在 NPU 上，Dato 在 GEMM 内核上实现了高达 84%的硬件利用率，在注意力内核上相对于最先进的商业框架实现了 2.81 倍的加速。
在 FPGA 上，Dato 生成的 定制脉动阵列达到了理论峰值性能的 98%。

作为未来工作，我们将把编译器从本地决策扩展到全局图级优化。我们还计划从任务规范自动生成高性能内核，减少对手写库的依赖。最后，我们将扩大目标范围，评估如 NVIDIA Hopper 及其具有 TMA 的后继者等加速器，以及其他可编程数据流结构。

随着人工智能计算的不断发展，数据流加速器将扮演越来越重要的角色。Dato 为这些复杂硬件的编程提供了全新的思路，有望成为下一代人工智能计算基础设施的关键组成部分。

Dato 的核心创新在于其类型系统的设计，将数据通信和数据分片都提升为一等类型， 这使得编译器能够在早期阶段进行深入的分析和优化，而无需依赖复杂且容易出错的手动调度。这种设计哲学不仅提高了性能，还显著降低了编程复杂性，为数据流加速器的广泛采用铺平了道路。

关注“鲸栖”小程序，掌握最新AI资讯

本文来自网络搜集，不代表鲸林向海立场，如有侵权，联系删除。转载请注明出处：http://www.itsolotime.com/archives/18296