# From Loop Nests to Silicon: Mapping AI Workloads onto AMD NPUs with MLIR-AIR

arXiv: https://arxiv.org/abs/2510.14871 · HTML: https://mlsanity.com/paper/2510.14871/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 44 references that are in the map.
- First posted on arXiv 16 Oct 2025.

## Builds on (most cited first, 12 of 12)

- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings](https://mlsanity.com/paper/2304.01433/) (2023)
- [SambaNova SN40L: Scaling the AI Memory Wall with Dataflow and Composition of Experts](https://mlsanity.com/paper/2405.07518/) (2024)
- [Efficiency, Expressivity, and Extensibility in a Close-to-Metal NPU Programming Interface](https://mlsanity.com/paper/2504.18430/) (2025)
- [NERO: A Near High-Bandwidth Memory Stencil Accelerator for Weather Prediction Modeling](https://mlsanity.com/paper/2009.08241/) (2020)
- [FPGA-Based Near-Memory Acceleration of Modern Data-Intensive Applications](https://mlsanity.com/paper/2106.06433/) (2021)
- [Revet: A Language and Compiler for Dataflow Threads](https://mlsanity.com/paper/2302.06124/) (2023)
- [Union: A Unified HW-SW Co-Design Ecosystem in MLIR for Evaluating Tensor Operations on Spatial Accelerators](https://mlsanity.com/paper/2109.07419/) (2021)
- [RUBICON: A Framework for Designing Efficient Deep Learning-Based Genomic Basecallers](https://mlsanity.com/paper/2211.03079/) (2022)
- [SPARTA: Spatial Acceleration for Efficient and Scalable Horizontal Diffusion Weather Stencil Computation](https://mlsanity.com/paper/2303.03509/) (2023)
- [Unlocking the AMD Neural Processing Unit for ML Training on the Client Using Bare-Metal-Programming Tools](https://mlsanity.com/paper/2504.03083/) (2025)
- [Designing, Modeling, and Optimizing Data-Intensive Computing Systems](https://mlsanity.com/paper/2208.08886/) (2022)

## Built on by (most cited first, 2 of 2)

- [From Human Guidance to Autonomy: Agent Skill System for End-to-End LLM Deployment on Spatial NPUs](https://mlsanity.com/paper/2606.07586/) (2026)
- [Reconfigurable Computing Challenge: Real-Time Graph Neural Networks for Online Event Selection in Big Science](https://mlsanity.com/paper/2605.10612/) (2026)

## Abstract (arXiv)

General-purpose compilers abstract away parallelism, locality, and synchronization, limiting their effectiveness on modern spatial architectures. As modern computing architectures increasingly rely on fine-grained control over data movement, execution order, and compute placement for performance, compiler infrastructure must provide explicit mechanisms for orchestrating compute and data to fully exploit such architectures. We introduce MLIR-AIR, a novel, open-source compiler stack built on MLIR that bridges the semantic gap between high-level workloads and fine-grained spatial architectures such as AMD's NPUs. MLIR-AIR defines the AIR dialect, which provides structured representations for asynchronous and hierarchical operations across compute and memory resources. AIR primitives allow the compiler to orchestrate spatial scheduling, distribute computation across hardware regions, and overlap communication with computation without relying on ad hoc runtime coordination or manual scheduling. We demonstrate MLIR-AIR's capabilities through two case studies: matrix multiplication and the multi-head attention block from the LLaMA 2 model. For matrix multiplication, MLIR-AIR achieves up to 78.7% compute efficiency and generates implementations with performance almost identical to state-of-the-art, hand-optimized matrix multiplication written using the lower-level, close-to-metal MLIR-AIE framework. For multi-head attention, we demonstrate that the AIR interface supports fused implementations using approximately 150 lines of code, enabling tractable expression of complex workloads with efficient mapping to spatial hardware. MLIR-AIR transforms high-level structured control flow into spatial programs that efficiently utilize the compute fabric and memory hierarchy of an NPU, leveraging asynchronous execution, tiling, and communication overlap through compiler-managed scheduling.