PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks

This paper presents PolyBlocks, a modular, MLIR-based compiler infrastructure that utilizes pass pipelines and analytical cost models to automatically generate high-performance code for AI chips, demonstrating competitive or superior performance against existing frameworks like Torch Inductor and XLA on NVIDIA GPUs.

The Gist

Imagine you are a master chef (the AI Framework like PyTorch or JAX) who wants to cook a massive, complex feast (an AI Model). You have a recipe book with high-level instructions: "Chop the onions," "Sauté the garlic," "Simmer the sauce."

However, your kitchen (the Hardware, like an NVIDIA GPU) is a high-tech, automated factory with robotic arms, conveyor belts, and specialized ovens. The factory doesn't speak "Chef"; it speaks "Machine Code."

Currently, to get your food cooked, you usually hire a middleman (existing compilers like Torch Inductor or XLA). This middleman looks at your recipe and says, "Okay, for the onions, I'll use this pre-made, factory-fresh onion-dicing machine (a Vendor Library like CuDNN). For the sauce, I'll use that pre-made blender (CuBLAS)."

This works well if the factory has the right pre-made machines for every single dish. But what if you invent a new dish? Or what if you want to combine the onion chopping and garlic sautéing into one super-efficient step to save time? The middleman often can't do that because they are stuck using the pre-made machines. They might even have to stop the conveyor belt, dump the onions into a bowl, walk them to the garlic station, and start again. This wastes time and energy.

Enter PolyBlocks: The Ultimate Kitchen Architect

PolyBlocks is a new, revolutionary compiler infrastructure designed to be the ultimate kitchen architect. Instead of relying on pre-made machines, PolyBlocks looks at your high-level recipe and builds the perfect, custom-made factory floor from scratch for that specific dish.

Here is how it works, using simple analogies:

1. The "Lego" Approach (Modular & Reusable)

Imagine PolyBlocks is a giant box of high-quality Lego bricks.

• The Problem: Old compilers were like custom-built wooden houses. If you wanted to build a house for a new type of land (a new AI chip), you had to chop down new trees and start from zero.
• The PolyBlocks Solution: PolyBlocks is a set of standardized, reusable Lego blocks. Whether you are building a house for an NVIDIA GPU, an AMD chip, or a future chip we haven't invented yet, you use the same core blocks. You just snap them together in a slightly different order. This makes it incredibly fast to build compilers for new hardware.

2. The "Smart Assembly Line" (Fusion)

In a normal kitchen, you might chop onions, put them in a bowl, walk to the stove, dump them in, then chop garlic, put them in a bowl, walk back, and dump them in.

• PolyBlocks' Magic: PolyBlocks realizes that the onions and garlic are going to the same pot. It rewrites the recipe so the chef chops the onions directly into the pan, then immediately chops the garlic into the same pan.
• The Result: This is called Fusion. It eliminates the "walking to the bowl" (moving data between memory and the processor). In AI terms, this saves massive amounts of time because the computer doesn't have to constantly fetch data from the slow "fridge" (Global Memory) and put it in the "countertop" (Fast Memory).

3. The "Tiling" Strategy (Organizing the Workspace)

Imagine you have a huge pile of 10,000 potatoes to peel.

• Old Way: You try to peel them all at once on a tiny table. You keep running out of space, so you have to keep moving potatoes back and forth.
• PolyBlocks' Way: PolyBlocks breaks the 10,000 potatoes into small, manageable piles of 100 (called Tiling). It peels one pile completely, clears the table, and moves to the next.
• Why it matters: This ensures that the potatoes currently being peeled are always right under your hands (in the fast on-chip memory), so you never have to stop to run to the storage room.

4. The "Specialized Robot" (Mapping to Matrix Units)

Modern AI chips have special "super-arms" (Matrix Units or Tensor Cores) designed to do math on grids of numbers incredibly fast.

• The Challenge: These super-arms only work if the ingredients are arranged in a very specific grid shape.
• PolyBlocks' Skill: PolyBlocks is like a master organizer that reshapes the ingredients (the data) into the perfect grid before handing them to the super-arm. It doesn't just say "Do math"; it says, "Here is the math, arranged exactly how your super-arm likes it, so it can run at 100% speed."

5. The "Attention" Trick (The Transformer Secret Sauce)

Modern AI (like Chatbots) uses a mechanism called "Attention" to focus on important words. This is notoriously slow and memory-heavy.

• The Old Way: Existing compilers often use a pre-written "Flash Attention" script. It's fast, but it's a black box. If you tweak the recipe slightly, the script might break or become slow.
• PolyBlocks' Way: PolyBlocks builds the "Attention" step from the ground up, automatically figuring out how to combine the math steps so they happen in one smooth motion, without ever stopping to save data to the slow memory. It's like a chef who knows exactly how to juggle all the ingredients so none of them ever hit the floor.

The Results: Why Should You Care?

The paper tested PolyBlocks against the current industry leaders (Torch Inductor and XLA) on NVIDIA GPUs.

• The Competition: The leaders are like a team using a mix of pre-made machines and custom tools. They are very good, but they are limited by the pre-made machines.
• PolyBlocks: PolyBlocks is like a team that builds its own custom tools on the fly.
• The Outcome: PolyBlocks matched or beat the leaders in many cases, even though the leaders were using the "best pre-made machines" available. For individual tasks (like matrix multiplication), PolyBlocks was just as fast as the best hand-written code from experts.

The Bottom Line

PolyBlocks is a compiler infrastructure that stops relying on "pre-made parts" and instead automatically designs the perfect factory floor for any AI model, on any chip.

It takes the messy, high-level code that data scientists write and transforms it into a hyper-efficient, custom-built machine code that runs as fast as humanly possible. It's the difference between hiring a contractor who uses standard blueprints and hiring an architect who designs a custom home specifically for your family's needs, built with the most efficient materials available.

This means that in the future, as we invent new, weird, and powerful AI chips, we won't have to wait years for software engineers to manually rewrite code for them. PolyBlocks can snap the Lego bricks together and get the new chip running instantly.

Technical Summary

Here is a detailed technical summary of the paper "PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks."

1. Problem Statement

The current landscape of AI compilation faces a significant gap between high-level programming frameworks (PyTorch, JAX, TensorFlow) and specialized hardware accelerators (GPUs, TPUs, custom AI chips).

• Reliance on Vendor Libraries: State-of-the-art compilers like Torch Inductor (for PyTorch) and XLA (for JAX/TensorFlow) achieve high performance primarily by relying on heavily optimized, hand-written vendor libraries (e.g., CuDNN, CuBLAS, FlashAttention) for compute-intensive operators. This limits portability to new hardware architectures that lack mature libraries.
• Limited Fusion: While these compilers support operator fusion, they often struggle to fuse complex, cross-operator patterns (like the entire Attention layer) or imperfectly nested loops because they rely on intermediate representations (IR) that are not optimized for deep, slicing-based transformations.
• Scalability and Portability: Building compilers for new AI chips often requires re-implementing complex optimizations from scratch or writing low-level kernels (CUDA, Triton), which is not scalable. There is a lack of a reusable, fully code-generating infrastructure that can automatically generate high-performance code for diverse hardware without manual intervention.

2. Methodology

The authors propose PolyBlocks, a modular, reusable, MLIR-based compiler infrastructure designed to bridge high-level frameworks to low-level hardware intrinsics through a fully code-generating approach.

Core Architecture

PolyBlocks is built on the MLIR (Multi-Level Intermediate Representation) infrastructure and organized into a five-stage pass pipeline:

1. S1 & S2 (Frontend/Buffering): Target-neutral stages that lower high-level tensor operations to buffer semantics (memref) and then to affine loop nests. This stage handles the transition from framework-specific IRs (Torch FX, HLO, StableHLO) to a unified affine representation.
2. S3 (Mid-Level Optimizer): The core engine where the majority of optimizations occur. It operates on affine and memref dialects. Key features include:
   • Lightweight Affine Analysis: Instead of heavy polyhedral integer-set operations, PolyBlocks uses efficient linear checks to analyze memory access patterns, enabling fast compile times even for models with thousands of loops.
   • Slicing-Based Fusion: A novel fusion technique that computes the specific "slice" of a producer needed by a consumer. This allows fusing imperfectly nested loops and pulling producers into consumers (or vice versa) to eliminate intermediate buffers, even if it introduces minor redundant computation (controlled by cost models).
   • Two-Phase Tiling & Fusion: The system tiles key nests (matmuls, convolutions) first, then performs fusion into these tiled structures, and finally tiles remaining nests for locality/parallelism.
3. S4 & S5 (Backend): Converts optimized affine nests into target-specific dialects (e.g., gpu, nvvm, llvm). This stage handles parallelism mapping, synchronization, and lowering to target intrinsics.

Key Technical Innovations

• Fully Code-Generating: PolyBlocks generates all low-level code (intrinsics) from scratch, avoiding reliance on external libraries. This ensures portability to new chips where libraries may not exist.
• Automatic Attention Layer Fusion: PolyBlocks automatically fuses the entire Attention layer (MatMul, Softmax, etc.) using a sequence of passes (reduce-reduce-fusion, wmma-fusion). It handles the complex data dependencies of the $Q \times K^T$ output and performs online softmax computation, eliminating global memory round-trips.
• On-the-Fly Packing for Convolutions: Instead of materializing large im2col matrices (which causes memory overhead), PolyBlocks transforms convolutions into matrix multiplications with on-the-fly packing. It loads and packs data into fast on-chip memory tiles dynamically during execution.
• Matrix Unit Mapping: The infrastructure automatically maps affine loops to hardware matrix units (e.g., NVIDIA Tensor Cores/WMMA) by analyzing access patterns and inserting the necessary vectorization and load/store instructions.

3. Key Contributions

1. Modular MLIR Infrastructure: A reusable compiler stack that separates target-agnostic optimizations (S1-S3) from target-specific lowering (S4-S5), enabling rapid porting to new AI chips.
2. Advanced Fusion Techniques: Implementation of slicing-based fusion that goes beyond traditional loop fusion, allowing the fusion of imperfectly nested loops and complex operator compositions (like Attention) that existing compilers cannot handle.
3. Automatic High-Performance Code Generation: Demonstrated ability to generate code for convolutions and matmuls that competes with hand-tuned vendor libraries (CuDNN, CuBLAS) and mid-level frameworks (Triton) without manual kernel writing.
4. Comprehensive Optimization Pipeline: Integration of multi-level tiling, on-chip scratchpad management, vectorization, and specialized attention fusion into a single automated flow.

4. Experimental Results

The authors evaluated PolyBlocks on NVIDIA A10 and A100 GPUs targeting PyTorch and JAX workloads.

• End-to-End Performance:
  • PyTorch (Batch Size 1): PolyBlocks is 2.15x faster than eager execution and 1.4x faster than Torch Inductor on average. It outperforms TensorRT by 2.4x.
  • PyTorch (Batch Size 8): PolyBlocks matches Torch Inductor performance (0.97x) and is 1.8x faster than eager execution.
  • JAX: PolyBlocks is 2.12x faster than JAX eager and 1.15x faster than XLA.
• Operator-Level Performance:
  • Convolutions: PolyBlocks-generated code is competitive with CuDNN across hundreds of convolution configurations. In nearly 50 cases, it outperformed CuDNN by >2x.
  • MatMul: Performance is on par with CuBLAS and hand-written Triton kernels.
  • Attention Layer: PolyBlocks significantly outperforms Inductor (which relies on FlashAttention kernels) in many configurations, achieving up to 0.69x (meaning Inductor is slower) or better speedups depending on the specific configuration, by fully fusing the layer and eliminating memory trips.
• Ablation Study:
  • Tensor Cores: Provide a geometric mean speedup of 17x over non-tensor-core code.
  • Cross-Operator Fusion: Provides a 2.87x speedup over code that is optimized individually but not fused.
  • Reduce-Reduce Fusion: Critical for attention layers, providing significant speedups (up to 29x in specific large models like Flux).

5. Significance

• Paradigm Shift: PolyBlocks challenges the industry standard of relying on vendor libraries for performance. It proves that a fully code-generating compiler can match or exceed the performance of hand-tuned libraries while offering superior portability.
• Future-Proofing for AI Chips: As the AI hardware landscape diversifies (beyond NVIDIA GPUs), the ability to automatically generate optimized code for new architectures without waiting for library vendors is crucial. PolyBlocks provides the infrastructure to do this.
• Scalability: By using lightweight affine analysis instead of full polyhedral integer-set computations, PolyBlocks achieves compile times in the order of seconds to tens of seconds, making it practical for large-scale models with thousands of operators.
• Unified Optimization: It unifies the optimization of diverse operators (convolutions, matmuls, attention) under a single framework, enabling cross-operator optimizations that were previously impossible or required manual intervention.

In conclusion, PolyBlocks represents a significant advancement in AI compiler infrastructure, demonstrating that high-level frameworks can be compiled to specialized hardware with high performance, full automation, and reusability, reducing the dependency on proprietary, hand-tuned libraries.