# PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 10 of its 58 references that are in the map.
- First posted on arXiv 6 Mar 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [cuDNN: Efficient Primitives for Deep Learning](https://mlsanity.com/paper/1410.0759/) (2014)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [SCENIC: A JAX Library for Computer Vision Research and Beyond](https://mlsanity.com/paper/2110.11403/) (2021)
- [ThunderKittens: Simple, Fast, and Adorable AI Kernels](https://mlsanity.com/paper/2410.20399/) (2024)
- [Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs](https://mlsanity.com/paper/2210.09603/) (2022)
- [Automatic Kernel Generation for Volta Tensor Cores](https://mlsanity.com/paper/2006.12645/) (2020)

## Abstract (arXiv)

We present the design and implementation of PolyBlocks, a modular and reusable MLIR-based compiler infrastructure for AI programming frameworks and AI chips. PolyBlocks is based on pass pipelines that compose transformations on loop nests and SSA, primarily relying on lightweight affine access analysis; the transformations are stitched together in specialized ways to realize high-performance code automatically by the use of analytical cost models and heuristics. The optimizations in these passes include multi-level tiling, fusion, on-chip scratchpad usage, mapping matmuls and convolutions to matrix units, fusing the attention layer, and several other transformations for parallelism and locality. They have been developed in a way that makes it easy to build PolyBlocks-based compilers to target new chips, reusing much of the infrastructure. PolyBlocks' design and architecture enable fully automatic code generation from high-level frameworks to low-level target-specific intrinsics. Experimental results from evaluating PolyBlocks-powered just-in-time compilation for PyTorch and JAX targeting NVIDIA GPUs show that it is able to match or outperform Torch Inductor and XLA in several cases, although the latter rely on a combination of vendor libraries and code generation. For individual operators like matmuls and convolutions, PolyBlocks-generated code is competitive with the best vendor-tuned libraries or hand-written kernels.