PolyBlocks: AI Compiler Infrastructure

• PolyBlocks is an MLIR-based compiler infrastructure that automatically translates high-level AI frameworks into low-level, target-specific code for AI chips.
• It employs a five-stage pass pipeline featuring affine loop transformations, slicing-based fusion, and efficient on-chip memory utilization for matrix operations.
• Empirical results demonstrate significant performance gains, with speedups over traditional frameworks and competitive results against tools like Torch Inductor and cuDNN.

Searching arXiv for recent and related papers on “PolyBlocks” and associated uses of the term. I searched arXiv for “PolyBlocks” and closely related entries, including the 2026 compiler paper, the 2022 magnetic block assembly environment, the 2026 polyform-enumeration paper, the 2023 topological interlocking paper, and the 2008 MAPEL optimization paper, which uses “polyblocks” in monotonic optimization. PolyBlocks most directly denotes an MLIR-based compiler infrastructure for AI programming frameworks and AI chips that is designed to be modular, reusable, and fully code-generating from high-level frameworks to low-level target-specific intrinsics (Bondhugula et al., 6 Mar 2026). The same string also has established technical uses in other research contexts: as a mathematical object in monotonic optimization, as a retroactive label applied to a magnetic-block assembly benchmark, and as a convenient descriptor for polyform-based block systems. The dominant contemporary usage in arXiv, however, is the 2026 compiler system, whose design centers on pass pipelines, loop nests, SSA, lightweight affine access analysis, analytical cost models, and heuristics for fully automatic performance-oriented code generation (Bondhugula et al., 6 Mar 2026).

1. Nomenclature and scope

The name PolyBlocks is not semantically uniform across the literature. In the compiler literature, it is the title of a specific system: “PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks” (Bondhugula et al., 6 Mar 2026). In monotonic optimization, polyblocks are finite unions of axis-aligned boxes anchored at the origin, used to outer-approximate normal feasible sets in the MAPEL algorithm for weighted throughput maximization (0805.2675). In embodied RL, the 2022 magnetic assembly environment is explicitly described as not using the name “PolyBlocks” in the original paper, although the supplied technical reading treats it as a PolyBlocks-style assembly domain (Ghasemipour et al., 2022). In combinatorial geometry, a “PolyBlocks-oriented” reading is used to interpret polyforms on tilings as catalogs of block systems, puzzles, and modular structures (Dobbelaere et al., 26 Feb 2026). In topological interlocking research, “PolyBlocks” is used as an interpretive frame for modular interlocking assemblies rather than as the paper’s formal name (Goertzen et al., 2023).

Usage	Meaning	Source
PolyBlocks	MLIR-based compiler infrastructure for AI programming frameworks and AI chips	(Bondhugula et al., 6 Mar 2026)
polyblocks	Finite unions of boxes used in monotonic optimization	(0805.2675)
PolyBlocks-style assembly domain	Magnetic-block RL benchmark, not originally named PolyBlocks	(Ghasemipour et al., 2022)
PolyBlocks-oriented reading	Enumeration of polyforms on tilings for toys, puzzles, and structures	(Dobbelaere et al., 26 Feb 2026)

This distribution of meanings suggests that the term functions both as a proper noun and as a family resemblance label. A plausible implication is that the compiler usage should be treated as the primary encyclopedic referent, with the other senses recorded as technically distinct homonyms.

2. Compiler-system definition and architectural goals

As a compiler system, PolyBlocks is built on MLIR and is intended to bridge high-level AI programming frameworks and specialized AI chips in a fully code-generating way (Bondhugula et al., 6 Mar 2026). Its stated design goals include full automation, fully code-generating execution without dependence on pre-tuned vendor libraries, pass-pipeline modularity, reuse across frameworks, and performance portability. The system is positioned against three recurrent problems: library dependence and lack of fusion, multi-framework, multi-target reuse, and the difficulty of scaling polyhedral-style optimizations to modern models with large numbers of loop nests (Bondhugula et al., 6 Mar 2026).

PolyBlocks integrates with the existing MLIR ecosystem rather than introducing a standalone IR. For PyTorch, the supplied description specifies a path through Torch Dynamo, FX graph capture, and torch-mlir lowering into MLIR torch / aten / linalg dialects. For JAX, the entry IR is stablehlo. For TensorFlow, PolyBlocks can plug into the TF MLIR pipeline at the mhlo level. The infrastructure also contains a small polyblocks dialect for special lowering hooks or hints, including polyblocks.matrix_cast, which is used to “view” memrefs as 2D vector or matrix types (Bondhugula et al., 6 Mar 2026).

The central architectural claim is reuse. The same core engine serves PyTorch, JAX, and TensorFlow once their graphs are lowered to MLIR dialects such as linalg, mhlo, or stablehlo (Bondhugula et al., 6 Mar 2026). Transformations are expressed in terms of affine loops, SSA, memref, vector, and target properties such as memory sizes, matrix shapes, and parallelism. This design is presented as making it easy to build PolyBlocks-based compilers to target new chips while reusing much of the infrastructure (Bondhugula et al., 6 Mar 2026).

3. Pass pipeline, IR organization, and analysis model

The compiler core is organized as a five-stage MLIR pass pipeline (S1–S5) driven by a Python compiler driver (Bondhugula et al., 6 Mar 2026). S1 lowers tensor-typed operations into memref-typed operations through bufferization. S2 lowers named operations such as linalg.matmul and convolution operators into explicit affine loop nests using the affine and memref dialects. S3 is the mid-level optimizer and the core of PolyBlocks. S4 lowers affine and memref constructs into scf and target-specific parallel dialects such as gpu or omp. S5 lowers the resulting representation into LLVM IR and target-specific intrinsics, including NVVM for NVIDIA GPUs (Bondhugula et al., 6 Mar 2026).

The middle-end relies primarily on lightweight affine access analysis over MLIR’s affine constructs, with polyhedral integer-set machinery used only when necessary (Bondhugula et al., 6 Mar 2026). This is a deliberate design choice: the system does not adopt full polyhedral scheduling or large-scale AST generation as its default mechanism. Instead, it performs transformation directly on loop IR, using affine maps to answer locality and dependence questions, and uses ISL only in relatively rare cases. The paper’s positioning is explicit: PolyBlocks aims for a balance between expressive power and compile-time scalability for large models (Bondhugula et al., 6 Mar 2026).

The internal optimization workload is substantial. The description specifies 50–70 passes per target in S3, operating primarily on affine and memref, with arith, math, func, and vector for scalar and vector operations (Bondhugula et al., 6 Mar 2026). Target information is injected through module attributes and pass configuration parameters. The high degree of decomposition into passes is not incidental; it is the mechanism by which more specialized behaviors such as attention fusion, matrix-unit mapping, and implicit-GEMM convolution are assembled from reusable transformations. This suggests a compiler architecture that is intentionally more infrastructural than monolithic.

4. Transformations: tiling, fusion, scratchpads, matrix units, and attention

The optimization repertoire described for PolyBlocks includes 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 (Bondhugula et al., 6 Mar 2026). Tiling is applied hierarchically to align with thread blocks, warps, threads, scratchpads, registers, and matrix-unit tile shapes. The compiler first tiles “destination” nests such as matmuls and convolutions, then fuses surrounding producer and consumer nests into the tiled destination via slicing-based affine fusion (Bondhugula et al., 6 Mar 2026).

The fusion mechanism is explicitly slicing-based affine fusion. Given a producer store, a consumer load, and a chosen fusion depth, PolyBlocks computes the slice of the producer needed to satisfy exactly the consumer tile or iteration space at that depth, inserts the slice loop into the consumer nest, rewrites the consumer to use the newly computed value or a small private buffer, and may erase the original producer if all uses are covered (Bondhugula et al., 6 Mar 2026). The paper notes that this can introduce redundant computation, controlled by heuristics that tolerate approximately 0–10% redundancy (Bondhugula et al., 6 Mar 2026). The stated advantage is generality: fusion is not limited to conventional elementwise cases and extends to imperfect nests and stencil-like computations.

For accelerators, on-chip memory is treated as a first-class optimization target. PolyBlocks generates fast buffers in scratchpad memory, emits copy-in and copy-out nests in affine form, and uses cost models to decide which memrefs merit fast buffers and which tile sizes fit the on-chip capacity (Bondhugula et al., 6 Mar 2026). An additional pass overlaps data movement with compute through affine loop shifting and maps staged copies to NVIDIA asynchronous copy instructions such as cp.async, again under heuristic control (Bondhugula et al., 6 Mar 2026).

Matrix-unit support is implemented through generic matricization. The compiler identifies loop dimensions corresponding to $M$, $N$, and $K$, introduces polyblocks.matrix_cast views, tiles nests to match MMA shapes, and emits matrix-load, MMA-compute, and matrix-store operations for tensor-core execution (Bondhugula et al., 6 Mar 2026). Convolutions are transformed into matrix multiplication form with on-the-fly packing, a strategy described as similar in spirit to CUTLASS implicit GEMM (Bondhugula et al., 6 Mar 2026). The same mapping machinery is used for standalone matmuls, convolutions, and matmuls embedded within larger fused kernels.

Attention is treated as a distinguished case. PolyBlocks uses reduce–reduce fusion and WMMA fusion to combine QK^T, bias or masking operations, softmax reductions, and the V multiplication into a single fused affine nest, avoiding materialization of the full $N \times N$ attention matrix in off-chip memory (Bondhugula et al., 6 Mar 2026). The description states that the generated kernels resemble manually written FlashAttention kernels, but are derived automatically from affine nests (Bondhugula et al., 6 Mar 2026). This is one of the strongest examples of PolyBlocks’ claim that specialized high-performance code can emerge from pass-pipeline composition rather than operator-specific library calls.

5. Performance characteristics and empirical evaluation

The experimental claims in the supplied description are extensive. For end-to-end PyTorch inference at batch size 1 on A10, PolyBlocks is reported as having a geometric mean speedup of 2.15× over PyTorch eager, being 1.4× faster than Torch Inductor, and 2.4× faster than TensorRT (Bondhugula et al., 6 Mar 2026). At batch size 8 on A100, the corresponding claims are a geometric mean speedup of 1.8× over eager and 0.97× relative to Torch Inductor, described as essentially on par (Bondhugula et al., 6 Mar 2026). For JAX workloads on A100, the reported geometric mean speedups are 2.12× over JAX eager and 1.15× faster than XLA (Bondhugula et al., 6 Mar 2026).

The operator-level results are similarly strong. For convolutions, PolyBlocks-generated kernels are described as competitive with cuDNN across many shapes, and in approximately 50 cases exceed cuDNN performance by more than 2× (Bondhugula et al., 6 Mar 2026). For GEMMs on A100, performance is reported as very close to cuBLAS and comparable or better than Triton across a range of sizes (Bondhugula et al., 6 Mar 2026). For attention, the comparison is more mixed: Torch Inductor maps attention layers to hand-written FlashAttention kernels, and PolyBlocks’ fused attention pipeline is described as competitive with state-of-the-art specialized kernels, with some cases slower at large sequence lengths and some realistic model configurations matching or exceeding Inductor after full fusion (Bondhugula et al., 6 Mar 2026).

The ablation study attributes a large portion of the system’s gains to both hardware mapping and high-level fusion. Enabling tensor cores relative to scalar FP32 yields a geometric mean speedup of approximately 17×. Enabling cross-operator fusion relative to no fusion yields a geometric mean speedup of approximately 2.87×. For attention-heavy workloads, reduce–reduce fusion + WMMA fusion adds up to approximately 1.4× over no reduce–reduce fusion and approximately 1.07× over reduce–reduce fusion alone at the full-model level (Bondhugula et al., 6 Mar 2026).

The evaluation also includes compiler-scope observations. The description states that about 90% of 205 PyTorch models from the PyTorch 2 evaluation suite compile and run (Bondhugula et al., 6 Mar 2026). This is significant because the paper explicitly attributes some failures to missing lowerings in torch-mlir or other frontends rather than to the core PolyBlocks engine (Bondhugula et al., 6 Mar 2026). A plausible implication is that the compiler’s practical reach is partly coupled to the maturity of the MLIR frontends it depends on.

6. Related meanings, mathematical homonyms, and broader research context

Outside the compiler paper, the closest formally defined technical homonym is the optimization-theoretic polyblock. MAPEL defines a polyblock as the union of finitely many boxes of the form $[0,\mathbf{v}_i]$ and uses these objects to approximate a normal feasible set in transformed SINR space for global power control (0805.2675). Proper vertices are the nondominated vertices of the polyblock, and MAPEL iteratively refines an outer polyblock approximation of the feasible region by projection and vertex replacement until an $\epsilon$-optimal solution is obtained (0805.2675). This is mathematically unrelated to the compiler infrastructure, but it is the oldest explicit arXiv use of the term in the supplied material.

The term also appears as a retrospective or interpretive label in physical assembly research. The 2022 magnetic-block environment is described in the supplied reading as a PolyBlocks-style assembly domain, although the paper itself calls it the magnetic block assembly environment or magnetic assembly domain rather than “PolyBlocks” (Ghasemipour et al., 2022). That environment is a 3D MuJoCo domain with a fixed set of 16 cuboid blocks of 6 different types, magnet-based attachment, blueprint-defined targets, graph-structured observations, PPO with GNN policies, and large-scale training up to 1–2.5 billion environment steps (Ghasemipour et al., 2022). The retroactive naming emphasizes blueprint-based assembly and structured relational reasoning rather than any formal nomenclature in the original paper.

In combinatorial geometry, “Building with Blocks: Enumerating Polyforms on Tilings” is not a paper about a system named PolyBlocks; instead, the supplied technical reading presents the work as a PolyBlocks-oriented framework for analyzing connected cell complexes on periodic tilings and honeycombs (Dobbelaere et al., 26 Feb 2026). There, an $n$-polyform is an equivalence class of induced connected subgraphs of the dual graph of a tiling, considered up to translations or the full automorphism group, with explicit counts for cubic, truncated-octahedral, tetrahedral–octahedral, rectified cubic, disphenoid tetrahedral, bisymmetric hendecahedral, and snub trihexagonal settings (Dobbelaere et al., 26 Feb 2026). This usage is descriptive rather than nominative.

A further adjacent usage appears in topological interlocking research, where “PolyBlocks” is again interpretive. The paper on topological interlocking assemblies studies how arrangement alone changes structural behavior for a fixed Versatile Block geometry, using Interlocking Flows as a combinatorial surrogate for load transfer and validating the predictions with finite element analysis (Goertzen et al., 2023). The connection to “PolyBlocks” lies in the general modular-block design frame rather than in a formal term defined by the authors (Goertzen et al., 2023).

Taken together, these senses show that “PolyBlocks” now spans at least four distinct technical domains: compiler infrastructure, monotonic optimization geometry, embodied block assembly, and modular polyform or interlocking block systems. The common thread is compositional structure, but the meanings are not interchangeable.

7. Limitations, extensibility, and prospective significance

The compiler paper identifies several limitations. Operator coverage is partly constrained by frontend maturity. Cost models are analytical and heuristic rather than globally optimal. Hardware feature coverage beyond currently supported GPUs and CPUs requires new S4 and S5 paths. Attention configuration cross-product is broad but not exhaustive. Compile time can become substantial for very large models: for Stable Diffusion XL1 refiner U-Net, compilation is described as taking around 10 minutes, with approximately 7 of those minutes in S3 (Bondhugula et al., 6 Mar 2026). These are not incidental caveats; they define the current operational boundary of the infrastructure.

At the same time, the design is explicitly retargetable. T1, S1, and S2 are described as almost completely reused across targets; S3 is heavily shared, with specialization controlled largely by target attributes such as scratchpad size, supported MMA shapes, and parallelism parameters; S4 and S5 are the principal target-specific layers (Bondhugula et al., 6 Mar 2026). This suggests that the main research claim of PolyBlocks is not only performance on NVIDIA GPUs, but the proposition that a loop-nest–centric, MLIR-based, largely reusable compiler stack can replace a substantial fraction of the framework-specific and library-dependent machinery that dominates current AI compilation practice.

In that sense, PolyBlocks occupies a specific position within the compiler landscape. Relative to TorchInductor and XLA, it emphasizes loop-nest–level transformation and lower library dependence. Relative to mid-level kernel DSLs such as Triton, Pallas, Helion, and ThunderKittens, it emphasizes whole-model automatic compilation rather than manual tile-level kernel authoring. Relative to polyhedral compilers, it retains polyhedral reasoning only where necessary and otherwise works directly on affine IR to preserve compile-time tractability (Bondhugula et al., 6 Mar 2026). This suggests that PolyBlocks is best understood not as a single optimization, but as an attempt to make a particular compiler architecture—pass-pipeline modular, affine-analysis driven, and fully code-generating—viable for both AI frameworks and future AI chips.