从循环嵌套到硅片：将 AI 工作负载映射到 AMD NPU 上的 MLIR-AIR
计算与语言
2025-10-17 v1 硬件体系结构
机器学习
摘要
通用编译器抽象化了并行性、局部性和同步机制，这限制了其在现代空间架构上的有效性。随着现代计算架构日益依赖对数据传输、执行顺序和计算放置的细粒度控制来提升性能，编译器基础设施必须提供显式机制来协调计算和数据，以充分发挥此类架构的潜力。我们引入 MLIR-AIR，这是一个基于 MLIR 的新型开源编译堆栈，旨在弥合高级工作负载与细粒度空间架构（如 AMD 的 NPU）之间的语义鸿沟。MLIR-AIR 定义了 AIR 方言，提供了跨计算和内存资源的异步和分层操作的结构化表示。AIR 原语允许编译器进行空间调度、跨硬件区域分配计算，并在不依赖 ad hoc 运行时协调或手动调度的情况下实现计算与通信的重叠。我们通过两个案例研究演示了 MLIR-AIR 的能力：矩阵乘法和来自 LLaMA 2 模型的多头注意力块。对于矩阵乘法，MLIR-AIR 实现了最高达 78.7% 的计算效率，并生成的实现性能几乎与基于更低层次、接近金属的 MLIR-AIE 框架编写的领先级手优化矩阵乘法相当。对于多头注意力，我们展示了 AIR 接口支持使用约150行代码的融合实现，使能够以可管理的方式表达复杂的工作负载并高效映射到空间硬件。MLIR-AIR 将高级结构化控制流转换为高效利用 NPU 计算 fabric 和内存层次结构的空间程序，利用异步执行、分块和通信重叠通过编译器管理的调度。
引用
@article{arxiv.2510.14871,
title = {From Loop Nests to Silicon: Mapping AI Workloads onto AMD NPUs with MLIR-AIR},
author = {Erwei Wang and Samuel Bayliss and Andra Bisca and Zachary Blair and Sangeeta Chowdhary and Kristof Denolf and Jeff Fifield and Brandon Freiberger and Erika Hunhoff and Phil James-Roxby and Jack Lo and Joseph Melber and Stephen Neuendorffer and Eddie Richter and Andre Rosti and Javier Setoain and Gagandeep Singh and Endri Taka and Pranathi Vasireddy and Zhewen Yu and Niansong Zhang and Jinming Zhuang},
journal= {arXiv preprint arXiv:2510.14871},
year = {2025}
}