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

• URL: http://arxiv.org/abs/2510.14871v1
• Date: Thu, 16 Oct 2025 16:49:05 GMT
• Title: From Loop Nests to Silicon: Mapping AI Workloads onto AMD NPUs with MLIR-AIR
• Authors: Erwei Wang, Samuel Bayliss, Andra Bisca, Zachary Blair, Sangeeta Chowdhary, Kristof Denolf, Jeff Fifield, Brandon Freiberger, Erika Hunhoff, Phil James-Roxby, Jack Lo, Joseph Melber, Stephen Neuendorffer, Eddie Richter, Andre Rosti, Javier Setoain, Gagandeep Singh, Endri Taka, Pranathi Vasireddy, Zhewen Yu, Niansong Zhang, Jinming Zhuang,
• Abstract summary: General-purpose compilers abstract away parallelism, locality, and synchronization, limiting their effectiveness on modern spatial architectures.<n>We introduce MLIR-AIR, a novel, open-source compiler stack built on MLIR.<n>We demonstrate MLIR-AIR's capabilities through two case studies: matrix multiplication and the multi-head attention block from the LLaMA 2 model.
• Score: 6.2957456904504525
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Shuttling Compiler for Trapped-Ion Quantum Computers Based on Large Language Models [27.588919587175543]
  Trapped-ion quantum computers rely on connectivity operations to establish connectivity between sub-registers within a quantum processing unit.<n>We propose a layout compilation strategy based on language increasing models to address this challenge.<n>We demonstrate that it provides a foundation for developing LLM-based compilers for branched-ion quantum computers.
  arXiv  Detail & Related papers  (2025-12-19T19:29:09Z)
• An LLVM-Based Optimization Pipeline for SPDZ [0.0]
  We implement a proof-of-concept LLVM-based optimization pipeline for the SPDZ protocol.<n>Our front end accepts a subset of C with lightweight privacy annotations and lowers it to LLVM IR.<n>Our back end performs data-flow and control-flow analysis on the optimized IR to drive a non-blocking runtime scheduler.
  arXiv  Detail & Related papers  (2025-12-11T20:53:35Z)
• Eliminating Multi-GPU Performance Taxes: A Systems Approach to Efficient Distributed LLMs [61.953548065938385]
  We introduce the ''Three Taxes'' (Bulk Synchronous, Inter- Kernel Data Locality, and Kernel Launch Overhead) as an analytical framework.<n>We propose moving beyond the rigid BSP model to address key inefficiencies in distributed GPU execution.<n>We observe a 10-20% speedup in end-to-end latency over BSP-based approaches.
  arXiv  Detail & Related papers  (2025-11-04T01:15:44Z)
• xLLM Technical Report [57.13120905321185]
  We introduce xLLM, an intelligent and efficient Large Language Model (LLM) inference framework.<n>xLLM builds a novel decoupled service-engine architecture.<n>xLLM-Engine co-optimizes system and algorithm designs to fully saturate computing resources.
  arXiv  Detail & Related papers  (2025-10-16T13:53:47Z)
• Scalable Engine and the Performance of Different LLM Models in a SLURM based HPC architecture [3.746889836344766]
  This work elaborates on a High performance computing architecture based on Simple Linux Utility for Resource Management (SLURM)<n> Dynamic resource scheduling and seamless integration of containerized have been leveraged to manage CPU, GPU, and memory efficiently in multi-node clusters.<n>The obtained results pave ways for significantly more efficient, responsive, and fault-tolerant LLM inference on large-scale HPC infrastructures.
  arXiv  Detail & Related papers  (2025-08-25T09:11:27Z)
• CCLSTM: Coupled Convolutional Long-Short Term Memory Network for Occupancy Flow Forecasting [0.0]
  We propose textbfCoupled Convolutional LSTM (CTM), a lightweight, end-to-end trainable architecture based solely on convolutional operations.<n>CTM achieves state-of-the-art performance on occupancy flow metrics and, as of this submission, ranks (textst) in all metrics on the 2024 Occupancy and Flow Prediction Challenge leaderboard.
  arXiv  Detail & Related papers  (2025-06-06T14:38:55Z)
• Splitwiser: Efficient LM inference with constrained resources [0.29260385019352086]
  Splitwiser is a methodology that splits the two phases of an LLM inference request onto the same GPU.<n>By eliminating the need to transfer data across devices, Splitwiser aims to minimize network-related overheads.<n>We implement our proposed multiprocessing design on two widely-used and independent LLM architectures: Huggingface and vLLM.
  arXiv  Detail & Related papers  (2025-04-21T00:21:08Z)
• FusionLLM: A Decentralized LLM Training System on Geo-distributed GPUs with Adaptive Compression [55.992528247880685]
  Decentralized training faces significant challenges regarding system design and efficiency. We present FusionLLM, a decentralized training system designed and implemented for training large deep neural networks (DNNs) We show that our system and method can achieve 1.45 - 9.39x speedup compared to baseline methods while ensuring convergence.
  arXiv  Detail & Related papers  (2024-10-16T16:13:19Z)
• EPS-MoE: Expert Pipeline Scheduler for Cost-Efficient MoE Inference [49.94169109038806]
  This paper introduces EPS-MoE, a novel expert pipeline scheduler for MoE that surpasses the existing parallelism schemes.<n>Our results demonstrate at most 52.4% improvement in prefill throughput compared to existing parallel inference methods.
  arXiv  Detail & Related papers  (2024-10-16T05:17:49Z)
• CLSA-CIM: A Cross-Layer Scheduling Approach for Computing-in-Memory Architectures [0.1747623282473278]
  We present CLSA-CIM, a cross-layer scheduling algorithm for tiled CIM architectures. We integrate CLSA-CIM with existing weight-mapping strategies and compare performance against state-of-the-art (SOTA) scheduling algorithms.
  arXiv  Detail & Related papers  (2024-01-15T13:35:21Z)
• In Situ Framework for Coupling Simulation and Machine Learning with Application to CFD [51.04126395480625]
  Recent years have seen many successful applications of machine learning (ML) to facilitate fluid dynamic computations. As simulations grow, generating new training datasets for traditional offline learning creates I/O and storage bottlenecks. This work offers a solution by simplifying this coupling and enabling in situ training and inference on heterogeneous clusters.
  arXiv  Detail & Related papers  (2023-06-22T14:07:54Z)
• Harnessing Deep Learning and HPC Kernels via High-Level Loop and Tensor Abstractions on CPU Architectures [67.47328776279204]
  This work introduces a framework to develop efficient, portable Deep Learning and High Performance Computing kernels. We decompose the kernel development in two steps: 1) Expressing the computational core using Processing Primitives (TPPs) and 2) Expressing the logical loops around TPPs in a high-level, declarative fashion. We demonstrate the efficacy of our approach using standalone kernels and end-to-end workloads that outperform state-of-the-art implementations on diverse CPU platforms.
  arXiv  Detail & Related papers  (2023-04-25T05:04:44Z)
• Energy-efficient Task Adaptation for NLP Edge Inference Leveraging Heterogeneous Memory Architectures [68.91874045918112]
  adapter-ALBERT is an efficient model optimization for maximal data reuse across different tasks. We demonstrate the advantage of mapping the model to a heterogeneous on-chip memory architecture by performing simulations on a validated NLP edge accelerator.
  arXiv  Detail & Related papers  (2023-03-25T14:40:59Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.