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

2025-10-16

AI Summary

Modern spatial architectures (e.g., AMD NPUs) demand fine-grained co-design of data movement, execution ordering, and computational layout—challenges that general-purpose compilers struggle to address due to their abstractions of parallelism, locality, and synchronization. This paper introduces MLIR-AIR, an open-source compiler stack built upon a structured asynchronous operation representation, featuring the AIR dialect to enable efficient mapping of AI workloads onto spatial hardware. Its key contributions include compiler-driven spatial scheduling, computational partitioning, and communication–computation overlap—eliminating reliance on manual scheduling or runtime coordination. Experimental evaluation demonstrates 78.7% of peak computational efficiency for matrix multiplication, matching hand-optimized code; moreover, a fused multi-head attention module achieves high performance with only ~150 lines of source code.

Technology Category

• Machine Learning: Hardware-aware ML
• Planning, Routing, and Scheduling: Planning with Language Models
• Computer Vision: Remote Sensing / Geospatial AI

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Economics, Online Markets and Human Computation: Architectures and workflows that use LLMs for crowd work
• Security and Privacy: Security and privacy of machine learning and AI applications

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.

Problem

• Bridging semantic gap between AI workloads and spatial architectures
• Providing explicit compiler mechanisms for spatial compute orchestration
• Enabling efficient mapping of complex workloads to NPU hardware

Innovation

• MLIR-AIR compiler stack bridges high-level workloads and NPUs
• AIR dialect enables asynchronous hierarchical compute and memory operations
• Compiler orchestrates spatial scheduling and communication-computation overlap