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.
Figures & tables
Figure 1. An indicative figure showing the trade-offs between high, mid, and low-level programming frameworks while delivering comparable performance. Matmul with PyTorch, Triton, and CUDA, all of which deliver close to peak performance. CUDA courtesy: Lei Mao ( Mao, 2024 ) .
Figure 2. PyTorch functions JIT compiled with PolyBlocks. Incremental programmability.
Figure 3. PolyBlocks compiler engine powering multiple programming frameworks and hardware. MLIR dialect lowering paths and how new targets are supported (XYZ in this figure).
Figure 4. Overall components and high-level architecture of a compiler built from PolyBlocks.
Figure 6. Affine slicing-based fusion. The first stencil nest is pulled into the second one by computing the producer slice needed by the consumer. Fusion happens into right under the 682-iteration %k loop of Figure 6 (a), i.e., given an %i, %j, %k, the slice of the producer needed.
Figure 7. On-the-fly tiled mapping for convolutions.
Figure 9. PyTorch performance with various backends: eager, Inductor, TensorRT, and PolyBlocks. Batch size 1 on the NVIDIA A10. Torch 2.5.1, torch-tensorrt-2.5.0, TensorRT 10.3.0.
Figure 10. PyTorch performance with various backends: eager, Inductor, TensorRT, and PolyBlocks. Batch size 8 on the NVIDIA A100. Torch 2.5.1, torch-tensorrt-2.5.0, TensorRT 10.3.0.
Workload
JAX execution time (ms)
Speedup with PolyBlocks
eager
XLA
PolyBlocks
over eager
over XLA
Resnet Batch 8
6.3
2.97
1.87
3.37
1.59
Unet Batch 8
84.6
60.1
69.1
1.22
0.87
Vit Batch 8
10.6
5.09
4.59
2.32
1.11
Table 1. JAX with XLA vs PolyBlocks. On NVIDIA A100, JAX 5.0 (Jan 2025) with CUDA 12.9, Flax 0.10.4.
Figure 11. Performance of convolution operator code generated fully with PolyBlocks on convolutions sourced from various AI models. cuDNN 9.10.2, PyTorch 2.9.0+cu128.
Figure 12. Performance comparison: CuBLAS, Triton, and PyTorch-PolyBlocks on a single matmul on NVIDIA A100. PyTorch 2.7.1. OpenAI Triton 3.3.1. CuBLAS 12.6.4.1.
Problem size (batch, num-heads, sequence length, dimensionality of head)
PolyBlocks
Inductor
PolyBlocks speedup over Inductor
No fusion
Only rr-fusion
rr+wmma fusion
No fusion
Only rr-fusion
rr+wmma fusion
B=1, H=16, N=512, D=64
0.0671
0.0516
0.0253
0.0174
0.26
0.34
0.69
B=1, H=16, N=2048, D=64
1.7821
0.4770
0.2002
0.1266
0.07
0.27
0.63
B=1, H=16, N=8192, D=64
487.12
6.2722
2.4359
1.5591
0.00
0.25
0.64
B=1, H=16, N=512, D=128
0.0854
0.1124
0.0417
0.0449
0.53
0.40
1.08
B=1, H=16, N=2048, D=128
1.8915
1.0164
0.3024
0.2296
0.12
0.23
0.76
Table 2. Attention layer performance for configurations drawn from various models. RR-fusion is reduce-reduce fusion. Torch 2.9.0, CUDA 12.8. NVIDIA A100. Execution times are in milliseconds.
Model
Execution times (milliseconds)
Speedup (with all opts) over
w/o tensor core
w/o fusion
w/o rr fusion
w/o wmma fusion
All opts
no-tensor-core
no-fusion
no-rr-fusion
no-wmma-fusion
AlexNet
2.11
0.35
0.28
0.28
0.28
7.57
1.24
1.00
1.00
ConvNext
12.32
2.83
0.94
0.93
0.93
13.20
3.03
1.01
1.00
DenseNet
9.14
5.31
2.26
2.27
2.23
4.11
2.39
1.02
1.02
EfficientNet2
11.82
4.87
1.42
1.42
1.40
8.43
3.47
1.01
1.01
GoogLeNet
4.19
2.02
0.57
0.57
0.55
7.65
3.68
1.03
1.04
Table 3. Torch-PolyBlocks performance impact of key optimizations isolated. On the A100.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 13. PyTorch i8 quantized performance with various backends: eager, Inductor, TensorRT, and PolyBlocks. NVIDIA A10. Torch 2.5.1, torch-tensorrt-2.5.0, TensorRT 10.3.0.
Figure 14. PyTorch i8 quantized performance with various backends: eager, Inductor, TensorRT, and PolyBlocks. NVIDIA A100. Torch 2.5.1, torch-tensorrt-2.5.0, TensorRT 10.3.0.
Figure 15. PyTorch-PolyBlocks i8 quantized performance vs fp16/fp32: batch size 8 on the NVIDIA A10.
Workload
IREE (milliseconds)
PolyBlocks (milliseconds)
fp32 with AMP
fp16
mixed fp16/fp32
Alexnet batch 8
17.92
15.21
0.36
VGG19 batch 8
225.59
189.65
2.81
Unet batch 8
1353.88
1303.28
13.2
DenseNet batch 8
Crash
Crash
3.49
Inception batch 8
Crash
Crash
1.85
Appendix
Table 4. Comparison with IREE. IREE version 3.8.0 (PIP). All execution times are in milliseconds.
Figure 17. PyTorch model coverage with PolyBlocks.