# AIE4ML: An End-to-End Framework for Compiling Neural Networks for the Next Generation of AMD AI Engines

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 30 references that are in the map.
- First posted on arXiv 17 Dec 2025.

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

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [MLP-Mixer: An all-MLP Architecture for Vision](https://mlsanity.com/paper/2105.01601/) (2021)
- [LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale](https://mlsanity.com/paper/2208.07339/) (2022)
- [hls4ml: An Open-Source Codesign Workflow to Empower Scientific Low-Power Machine Learning Devices](https://mlsanity.com/paper/2103.05579/) (2021)
- [Compressing deep neural networks on FPGAs to binary and ternary precision with HLS4ML](https://mlsanity.com/paper/2003.06308/) (2020)
- [AdaPT: Fast Emulation of Approximate DNN Accelerators in PyTorch](https://mlsanity.com/paper/2203.04071/) (2022)
- [Fast Jet Tagging with MLP-Mixers on FPGAs](https://mlsanity.com/paper/2503.03103/) (2025)
- [Efficiency, Expressivity, and Extensibility in a Close-to-Metal NPU Programming Interface](https://mlsanity.com/paper/2504.18430/) (2025)
- [hls4ml: A Flexible, Open-Source Platform for Deep Learning Acceleration on Reconfigurable Hardware](https://mlsanity.com/paper/2512.01463/) (2025)
- [MaxEVA: Maximizing the Efficiency of Matrix Multiplication on Versal AI Engine](https://mlsanity.com/paper/2311.04980/) (2023)
- [TransAxx: Efficient Transformers with Approximate Computing](https://mlsanity.com/paper/2402.07545/) (2024)
- [GAMA: High-Performance GEMM Acceleration on AMD Versal ML-Optimized AI Engines](https://mlsanity.com/paper/2504.09688/) (2025)
- [AutoMM: Energy-Efficient Multi-Data-Type Matrix Multiply Design on Heterogeneous Programmable System-on-Chip](https://mlsanity.com/paper/2305.18698/) (2023)

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

- [Taming the Exponential: A Fast Softmax Surrogate for Integer-Native Edge Inference](https://mlsanity.com/paper/2604.02292/) (2026)
- [Design Rules for Extreme-Edge Scientific Computing on AI Engines](https://mlsanity.com/paper/2604.19106/) (2026)
- [Tempus: A Temporally Scalable Resource-Invariant GEMM Streaming Framework for Versal AI Edge](https://mlsanity.com/paper/2605.00536/) (2026)
- [Reconfigurable Computing Challenge: Real-Time Graph Neural Networks for Online Event Selection in Big Science](https://mlsanity.com/paper/2605.10612/) (2026)
- [TileFuse: A Fused Mixed-Precision Kernel Library for Efficient Quantized LLM Inference on AMD NPUs](https://mlsanity.com/paper/2606.11357/) (2026)

## Abstract (arXiv)

Efficient AI inference on AMD's Versal AI Engine (AIE) is challenging due to tightly coupled VLIW execution, explicit datapaths, and local memory management. Prior work focused on first-generation AIE kernel optimizations, without tackling full neural network execution across the 2D array. In this work, we present AIE4ML, the first comprehensive framework for converting AI models automatically into optimized firmware targeting the AIE-ML generation devices, also with forward compatibility for the newer AIE-MLv2 architecture. At the single-kernel level, we attain performance close to the architectural peak. At the graph and system levels, we provide a structured parallelization method that can scale across the 2D AIE-ML fabric and exploit its dedicated memory tiles to stay entirely on-chip throughout the model execution. As a demonstration, we designed a generalized and highly efficient linear-layer implementation with intrinsic support for fused bias addition and ReLU activation. Also, as our framework necessitates the generation of multi-layer implementations, our approach systematically derives deterministic, compact, and topology-optimized placements tailored to the physical 2D grid of the device through a novel graph placement and search algorithm. Finally, the framework seamlessly accepts quantized models imported from high-level tools such as hls4ml or PyTorch while preserving bit-exactness. In layer scaling benchmarks, we achieve up to 98.6% efficiency relative to the single-kernel baseline, utilizing 296 of 304 AIE tiles (97.4%) of the device with entirely on-chip data movement. With evaluations across real-world model topologies, we demonstrate that AIE4ML delivers GPU-class throughput under microsecond latency constraints, making it a practical companion for ultra-low-latency environments such as trigger systems in particle physics experiments.