PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks
Uday Bondhugula, Akshay Baviskar, Navdeep Katel, Vimal Patel, Anoop JS, and Arnab Dutta

TL;DR
PolyBlocks is a modular compiler infrastructure based on MLIR that automates high-performance code generation for AI chips and frameworks, leveraging analytical models and heuristics for optimization.
Contribution
It introduces a reusable, MLIR-based compiler framework that simplifies building target-specific AI chip compilers with automatic code generation capabilities.
Findings
PolyBlocks matches or outperforms Torch Inductor and XLA in several cases.
Generated code for key operators is competitive with vendor libraries.
The infrastructure facilitates targeting new chips with minimal effort.
Abstract
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…
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsParallel Computing and Optimization Techniques · Embedded Systems Design Techniques · Logic, programming, and type systems
