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

TL;DR

The paper tackles the gap between high-level AI workloads and spatial accelerator hardware by introducing MLIR-AIR, an MLIR-based compiler framework and AIR dialect that expose explicit spatial scheduling, dataflow, and synchronization. By representing asynchronous and hierarchical compute with air.launch, air.segment, air.herd, air.channel, and air.token, MLIR-AIR enables compiler-managed overlap of communication and computation and supports tile-based mappings to AMD NPUs. Through matmul and LLaMA 2 MHA case studies, it demonstrates competitive QoR, with up to 78.7% of peak throughput for certain data types and a 2.24× latency reduction for fused MHA kernels, while maintaining a compact, scalable IR. The approach shows promise for cross-platform spatial compilation, offering an open, modular path from high-level frontends to vendor-specific backends and runtimes, enabling portable yet high-performance AI acceleration.

Abstract

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.

Figures (14)

• Figure 1: AMD NPU architecture.
• Figure 2: MLIR-AIR stack overview.
• Figure 3: Impact of tiling strategy on data movement schedule in an output-stationary matrix multiplication. See Listing \ref{['lst:matmul_loop_nest']} for its schedule described using loop nests, where the for_all loops ii and jj were mapped to horizontal and vertical directions in the two-dimensional air.herd operations, respectively.
• Figure 4: Visualizations of ACDGs in loop iterations, including (\ref{['fig:scf_for_loop_carried']}) sequentialized for loop, (\ref{['fig:scf_for_race']}) multi-token for loop, and (\ref{['fig:scf_par_loop_reduced']}) parallel loop, and their respective MLIR-AIR specification. A circle represents an air.token, and a polygon represents a group of MLIR operations in the loop body. Listings (\ref{['fig:scf_for_loop_carried_code']}---\ref{['fig:scf_par_loop_reduced_code']}) demonstrates how each ACDG is represented inline of MLIR code.
• Figure 5: Visualizations of ACDGs in loop iterations, including (\ref{['fig:scf_for_loop_carried']}) sequentialized for loop, (\ref{['fig:scf_for_race']}) multi-token for loop, and (\ref{['fig:scf_par_loop_reduced']}) parallel loop, and their respective MLIR-AIR specification.