MLIR编译器优化揭秘：双缓冲与多线程如何实现AI核函数延迟隐藏与并行加速

关键词： MLIR、AI Kernel 编译器、延迟隐藏、并行性、双缓冲

• Analyzing Latency Hiding and Parallelism in an MLIR-based AI Kernel Compiler
• https://arxiv.org/pdf/2602.20204
• 6000 字，阅读 24 分钟

面向边缘设备的 AI 核函数编译，其性能优劣除受其他因素影响外，核心取决于编译器在分层内存与显式数据搬移的硬件特性下，挖掘并行性并隐藏内存延迟的能力。

本文提出了一套基准测试方法，并给出了基于 MLIR 的编译流水线中三种编译器可控优化机制的测试结果，分别为：向量化（Vec）、跨硬件上下文的多线程（MT）、以及利用乒乓式暂存缓冲区（两块独立的高速片上暂存区，交替用于数据接收和计算）实现DMA 传输与计算重叠的双缓冲（DB, double buffer）。

本文基于 Triton/Inductor 框架生成核函数，构建了消融梯度策略——该策略会逐步叠加优化手段以量化单一手段性能贡献的测试方法，以此区分向量化、多线程、双缓冲各自的性能贡献；同时以 GELU 作为代表性激活核函数，量化分析了多线程的加速比随问题规模的缩放规律。

实验结果表明：

1. 向量化为带宽敏感型核函数（执行速度主要受内存数据传输带宽限制的计算核）带来核心性能收益；
2. 当调度开销（线程创建、分发、同步等带来的性能损耗）被平摊后，多线程能实现显著的性能提升；
3. 而当数据传输与计算操作可实现重叠时（即核函数并非处于纯内存受限或纯计算受限的极端场景），双缓冲能进一步带来性能增益。

注：
* 内存受限：程序性能由内存访问速度/带宽决定，计算单元常因等待数据而空闲。
* 计算受限：程序性能由计算单元的运算速度决定，内存传输能及时满足计算需求。

本文目录

• 一、引言
• 二、实现细节：多线程与双缓冲
  • 2.1 多线程（MT）
  • 2.2 双缓冲（DB）
• 三、基准测试设置与方法
• 四、结果
  • 4.1 二维向量加法的消融梯度
  • 4.2 GELU：单线程与多线程的扩展性
  • 4.3 讨论
• 五、结论
• 参考文献

一、引言

核函数优化的自动化实现难度极高：手写核函数的性能至今难以被自动生成的代码超越，但实际生产系统又要求编译器能在快速迭代的硬件架构上，生成可移植且易维护的代码。

在边缘神经网络处理器（NPU）中，程序性能由三大因素决定：分层内存的硬件特性、DMA 管控的显式数据传输、以及让计算单元在数据传输过程中持续工作的任务调度策略。

端到端的模型性能结果虽有实际参考价值，但这类结果往往无法清晰体现出这种映射机制，即将算法逻辑映射到硬件执行的具体方式时，各种策略对性能的提升作用及背后原因。因此：

• 本文采用了一套可复现、面向核函数的评估方法，该方法与 Triton/Inductor 风格的代码生成逻辑保持一致；
• 通过在受控问题规模（固定且可复现的计算/数据规模）下对典型算子核函数进行基准测试，分别量化向量化（Vec）、多线程（MT）、双缓冲（DB）三种优化机制的性能影响[1]。
  • 其中，向量化挖掘数据级并行（对大量同类型数据执行相同操作实现的并行性）；
  • 多线程通过将独立计算分块（无数据依赖、可并行执行的计算子块）分发到不同硬件上下文（硬件为线程分配的独立执行环境，含寄存器、状态标志等），挖掘循环与区域级并行；
  • 双缓冲则通过让内存传输与计算操作重叠执行，减少计算单元的停滞时间（计算单元因等待数据而空闲的时间）。

为将性能收益精准归因于特定优化机制，本文通过一套简洁的消融梯度呈现测试结果，梯度层级如下：

标量执行 → 向量化 → 向量化+多线程 → 向量化+多线程+双缓冲

在该消融梯度中：

• 向量化层级单独验证SIMD 风格 lowering的性能收益；
• 向量化+多线程层级量化增量式的线程级加速比；
• 向量化+多线程+双缓冲层级则用于评估：在向量化与多线程优化已生效的基础上，显式延迟隐藏调度是否能带来额外的性能提升。

除呈现测试结果外，本文还详细描述了实现多线程与双缓冲所采用的具体 MLIR中间表示（IR）模式与编译 Pass 结构，让本研究的方法易于复现与扩展。

二、实现细节：多线程与双缓冲

本文的优化机制基于 MLIR 框架[2]实现，并遵循一个核心设计原则：尽可能长时间保留结构化中间表示中表达的计算意图，仅当编译器完成任务调度后，再将其 lowering 为运行时结构（程序执行阶段才生效的代码结构，如线程创建、同步原语）。

该原则能提升代码的可移植性，同时让编译器的转换优化操作更易验证。

2.1 多线程（MT）

硬件模型

本文的目标边缘 NPU 支持多线程向量执行，每个硬件线程都对应一个独立的向量上下文（如向量寄存器、谓词状态，其中向量寄存器用于存放批量数据，谓词状态用于控制向量指令的条件执行）。

若分块迭代空间（将大的计算迭代空间拆分为小的分块）可被无跨线程数据依赖（不同线程间的计算需要彼此的结果，导致无法并行）地拆分，该硬件特性可实现同一核函数的多个独立计算分块的并发执行。

两阶段多线程 lowering

本文将多线程的下推实现为两阶段流水线：第一阶段保留结构化并行性，第二阶段引入显式的分叉-合并（一种并行执行模型，先将任务拆分到多个线程并行执行，再合并结果完成同步）。

2.1 多线程（MT）

多线程优化旨在将计算分块映射到多个硬件线程上并行执行，以隐藏延迟并提升吞吐率。MLIR编译器通过两个结构化的阶段实现这一目标。

1. 虚拟线程构建阶段：此阶段将分块后的核函数（例如linalg.generic操作）通过scf.forall并行循环指令，重写为显式的并行执行形式。编译器会基于计算分块空间的维度，采用基于规模的收益启发式算法来判断并行化是否带来性能收益，并选择合适的任务分配策略（如块式或块循环式分配），以在迭代范围非均匀时实现线程间的负载均衡。

2. 异步线程构建阶段：此阶段利用MLIR的异步方言，将scf.forall指令转换为分叉-合并的表示形式。每个计算分块被封装为一个async.execute异步执行区域，该区域执行完成后会产生一个令牌，作为任务完成的标识。所有令牌被收集到一个异步任务组中。在后续的依赖计算执行前，通过async.await all指令实现屏障同步，确保所有并行任务完成。

中间表示中的标准分叉-合并框架

编译器生成的分叉-合并代码模式被设计为轻量且规整，以便直接下推至协程/任务运行时。其核心结构如下：

mlir
%group = async.create_group %N
scf.for %tile = ... {
%tok = async.execute { /* 分块计算逻辑 */ async.yield }
async.add_to_group %tok, %group
}
async.await_all %group

选择异步方言的原因

将多线程以结构化的分叉-合并形式保留在异步方言中，能够在延迟下推前，以声明式形式保留并行语义。这既便于进行编译器优化，又能被直接转换为运行时的任务调度逻辑。

2.2 双缓冲（DB）

双缓冲是一种软件流水线策略，通过交替使用两块暂存缓冲区（乒乓区），实现数据传输与计算操作的重叠执行，从而隐藏数据移动的延迟。

该技术与模调度和流水线循环执行密切相关。本文将双缓冲的实现分为两个解耦的阶段：首先构建结构化的流水线调度，然后集成面向硬件的异步传输原语。

阶段 1：结构化流水线构建

此阶段会匹配由分块操作生成的单缓冲分块循环标准形式，其核心指令序列为：memref.subview（内存子视图切分）→ memref.alloc（内存分配）→ memref.copy（内存拷贝），随后紧跟计算与写回操作。

基于此结构，编译器会构建显式的乒乓式调度：
* 生成序段代码，将第一个计算分块预取至乒乓缓冲区的“乒”区。
* 将主循环重构为两个交替执行的子核函数。
* 每个子核函数依次执行三个操作：
1. 将下一个计算分块预取至另一块缓冲区（“乓”区或“乒”区）。
2. 利用当前缓冲区中的数据执行计算。
3. 基于当前循环归纳变量重构内存子视图，完成计算结果的写回。
编译器通过一个布尔开关实现每次迭代的缓冲区切换，并添加轻量级属性作为锚点，供后续阶段精准识别预取、计算和写回区域。

阶段 2：异步 DMA 集成

此阶段将同步的内存拷贝操作替换为显式的异步 DMA 操作：
* 预取路径的代码被重写为带有唯一缓冲区标记的memref.dma_start指令。
* 在计算操作前插入memref.dma_wait指令，确保计算所需数据已驻留在紧耦合内存（TCM）中。
数据写回的拷贝操作也进行相同处理，同时生成均衡的标记释放逻辑，避免 DMA 标记资源泄漏。

这种分阶段转换确保了先确立正确的执行调度，再映射到硬件传输接口，符合经典的编译器设计实践。

双缓冲与多线程的组合实现

双缓冲与多线程可自然组合优化。当memref.dma_wait确保数据就绪后，计算区域可通过scf.forall虚拟线程或其 lowering 后的分叉-合并形式，实现多个子计算分块的并发执行。本文基准测试中的 向量化+多线程+双缓冲 层级即采用了该组合优化方式。

三、基准测试设置与方法

本文选取两个代表性核函数进行基准测试：
* 二维向量加法微基准测试：计算规模为 $2^{24}$ 个元素，适用于分析高内存开销场景下，向量化与延迟隐藏机制的相对性能影响。
* GELU 激活核函数：作为Transformer推理子图的典型算子，采用基于Triton框架实现的版本进行测试。

测试方法如下：
* 对于二维向量加法，评估消融梯度中的四个层级：标量执行、向量化、向量化+多线程、向量化+多线程+双缓冲。
* 对于GELU核函数，在一系列不同计算规模下测试，分别给出单线程与多线程执行的端到端延迟，以展现多线程调度开销的平摊效果与性能缩放规律。

所有延迟结果以微秒（µs）为单位，加速比以 单线程耗时 / 多线程耗时 的比值表示。

四、结果

4.1 二维向量加法的消融梯度

图1总结了二维向量加法核函数在消融梯度各层级的性能表现。

图1：二维向量加法（$2^{24}$个元素）的消融梯度（数值越低性能越好）。层级分别为：标量执行（132479 µs）、向量化（3210 µs）、向量化+多线程（3000 µs）、向量化+多线程+双缓冲（2689 µs）。

• 向量化带来了主导性的性能提升，延迟从132479 µs降至3210 µs，加速比约41.3倍。这与带宽敏感型核函数能从SIMD执行中直接获益的特性一致。
• 多线程与双缓冲带来了幅度更小但可量化的增量性能提升，延迟从3210 µs逐步降至2689 µs。

这表明，在向量化优化生效后，剩余的性能提升主要源于同步开销的降低与传输-计算操作的部分重叠效应，而非算术吞吐率的提升。

4.2 GELU：单线程与多线程的扩展性

图2呈现了GELU核函数在不同问题规模下的延迟表现。

图2：GELU核函数在单线程与多线程执行下，延迟随问题规模（元素数量为$2^N$尺度）的变化。

图3则给出了对应的多线程加速比。

图3：GELU核函数的多线程加速比（单线程耗时/多线程耗时）随问题规模（元素数量为$2^N$尺度）的变化。

实验结果表明，多线程在所有测试的问题规模下均能提升性能，且加速比随问题规模的增大而提升：当计算规模为 1048576 个元素时，加速比达到约 3.91 倍（单线程延迟 12947 µs，多线程延迟 3313 µs）。

该趋势表明，随着每个线程的计算量增加，分叉-合并调度带来的固定开销被不断平摊；而大问题规模下加速比趋于饱和的现象，则表明此时出现了 共享性能瓶颈（多个线程竞争导致的性能瓶颈），如内存带宽受限、屏障同步开销增大等。

4.3 讨论

在本文测试的两类核函数中，向量化均带来了核心的 一阶性能收益（最主要、最显著的性能提升），这一特性在带宽敏感型场景中尤为突出。

当存在 充足的并行余量（存在足够多可被拆分并并行执行的计算任务）、且核函数的计算规模足够大以平摊调度开销时，多线程能带来显著的增量性能提升。

而当数据传输与计算操作的耗时均占据显著比例时，双缓冲能带来额外的性能增益；在纯内存受限的极端场景下，双缓冲的性能提升会受传输带宽的限制；在纯计算受限的极端场景下，传输与计算的重叠机会则会大幅减少。

五、结论

本文提出了一套简洁、可复现的方法，用于分析基于 MLIR 的 AI 核编译器中的向量化、多线程与双缓冲优化机制。

本文设计的消融梯度能让研究者将性能提升精准归因于特定的编译器优化机制，而非仅呈现端到端的加速比结果。实验测量结果表明：

• 向量化是性能提升的基础优化手段；
• 当调度开销被平摊后，多线程能实现大幅的性能提升；
• 而当数据传输与计算操作可实现重叠时，双缓冲能带来增量的性能收益。

未来的研究工作将进一步扩充基准测试的核函数集，如 RMSNorm 均方根归一化、softmax 软最大化函数，并将本文的测量结果与预测模型关联，建立 重叠效率（传输与计算重叠的时间占比）与 DMA 吞吐率、暂存缓冲区容量之间的量化关系。

参考文献

• [1] Philippe Tillet, H. T. Kung, and David Cox. Triton: An intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages (MAPL), 2019.
• [2] Chris Lattner, Mehdi Amini, Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasilache, and Oleksandr Zinenko. Mlir: A compiler infrastructure for the end of moore’s law, 2020.
• [3] LLVM Project. Mlir async dialect. https://mlir.llvm.org/docs/Dialects/ AsyncDialect/, 2026. Accessed 2026-01-21.
• [4] Monica S. Lam. Software pipelining: An effective scheduling technique for vliw machines. In Proceedings of the ACM SIGPLAN 1988 Conference on Programming Language Design and Implementation (PLDI), 1988.
• [5] B. Ramakrishna Rau. Iterative modulo scheduling: An algorithm for software pipelining loops. Proceedings of the 27th Annual International Symposium on Microarchitecture (MICRO), 1994.

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