# Towards a high-performance AI compiler with upstream MLIR

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 3 of its 7 references that are in the map.
- First posted on arXiv 15 Apr 2024.

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

- [Tensor Processing Primitives: A Programming Abstraction for Efficiency and Portability in Deep Learning & HPC Workloads](https://mlsanity.com/paper/2104.05755/) (2021)
- [oneDNN Graph Compiler: A Hybrid Approach for High-Performance Deep Learning Compilation](https://mlsanity.com/paper/2301.01333/) (2023)
- [Harnessing Deep Learning and HPC Kernels via High-Level Loop and Tensor Abstractions on CPU Architectures](https://mlsanity.com/paper/2304.12576/) (2023)

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

- [Hexcute: A Compiler Framework for Automating Layout Synthesis in GPU Programs](https://mlsanity.com/paper/2504.16214/) (2025)
- [Space Filling Curves is All You Need: Communication-Avoiding Matrix Multiplication Made Simple](https://mlsanity.com/paper/2601.16294/) (2026)
- [RISC-V and machine learning: a survey](https://mlsanity.com/paper/2609.20677/) (2026)

## Abstract (arXiv)

This work proposes a compilation flow using open-source compiler passes to build a framework to achieve ninja performance from a generic linear algebra high-level abstraction. We demonstrate this flow with a proof-of-concept MLIR project that uses input IR in Linalg-on-Tensor from TensorFlow and PyTorch, performs cache-level optimizations and lowering to micro-kernels for efficient vectorization, achieving over 90% of the performance of ninja-written equivalent programs. The contributions of this work include: (1) Packing primitives on the tensor dialect and passes for cache-aware distribution of tensors (single and multi-core) and type-aware instructions (VNNI, BFDOT, BFMMLA), including propagation of shapes across the entire function; (2) A linear algebra pipeline, including tile, fuse and bufferization strategies to get model-level IR into hardware friendly tile calls; (3) A mechanism for micro-kernel lowering to an open source library that supports various CPUs.