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.
In this paper, we present Hexagon-MLIR,an open-source compilation stack that targets Qualcomm Hexagon Neural Processing Unit (NPU) and provides unified support for lowering Triton kernels and PyTorch models . Built using the MLIR framework, our compiler applies a structured sequence of passes to exploit NPU architectural features to accelerate AI workloads. It enables faster deployment of new Triton kernels (hand-written or subgraphs from PyTorch 2.0), for our target by providing automated compilation from kernel to binary. By ingesting Triton kernels, we generate mega-kernels that maximize data locality in the NPU's Tightly Coupled Memory (TCM), reducing the bandwidth bottlenecks inherent in library-based approaches. This initiative complements our commercial toolchains by providing developers with an open-source MLIR-based compilation stack that gives them a path to advance AI compilation capabilities through a more flexible approach. Hexagon-MLIR is a work-in-progress, and we are continuing to add many more optimizations and capabilities in this effort.
Mohammed Javed Absar, Muthu Baskaran, Abhikrant Sharma +22
Qualcomm Technologies International, Ltd. · Qualcomm Technologies, Inc. · Qualcomm India Private Limited
Compiler backends are expensive to build and maintain as programming models, workloads, and accelerators evolve. We investigate whether large language models can replace the conventional optimizing and lowering pipeline, a process that we call AI lowering. We study AI lowering from Triton to NVIDIA PTX: an LLM agent translates Triton kernels directly into PTX. We build an environment that evaluates candidate PTX, and an agentic harness in which an LLM translates Triton kernels into PTX. Across twelve common kernels on Ada, Hopper, and Blackwell GPUs and ten kernels from recent ML papers, AI lowering achieves 0.83x-3.34x the performance of autotuned Triton. The largest gains come from transformations that Triton's lowering pipeline does not perform, such as decoding packed binary weights directly into Tensor Core operands (3.34x on BitDelta), assigning each thread a complete softmax row in tensor memory (1.37x on FlashAttention), and reusing overlapping convolution windows (up to 2.23x). These results rely on a robust evaluation harness with comprehensive verification support. We build on Volta, an existing PTX verifier, and substantially extend it to support modern GPU architectures by introducing support for Blackwell's tcgen05 Tensor Core interface. This requires modeling three architectural features: managed tensor memory, descriptor-based operand layouts, and asynchronous execution coordinated through commits, waits, memory barriers, and proxy fences. We discuss the challenges involved in formalizing them, as well as the current limitations. Our results suggest an emerging future in which AI compilers replace custom-written intermediate representations and checkers, reducing the time and engineering effort required to bring up software for new general-purpose and custom chips.
We present Forge-UGC (FX Optimization and Register-Graph Engine for Universal Graph Compilation), a four-phase compiler for transformer deployment on heterogeneous accelerator hardware, validated on Intel AI Boost NPU. Existing frameworks such as OpenVINO and ONNX Runtime often use opaque compilation pipelines, limited pass-level visibility, and weak buffer management, which can lead to higher compilation cost and runtime overhead. Forge-UGC addresses this with a hardware-agnostic design that separates graph capture, optimization, intermediate representation lowering, and backend scheduling. Phase 1 captures graphs with torch.export at the ATen operator level, supporting modern transformer components such as rotary position embeddings, grouped-query attention, and SwiGLU without manual decomposition. Phase 2 applies six optimization passes: dead code elimination, common subexpression elimination, constant folding, attention fusion, operator fusion, and layout optimization, reducing graph node count by 14.2 to 21.9%. Phase 3 lowers the optimized graph into a typed intermediate representation with explicit virtual register assignments. Phase 4 performs liveness analysis, linear-scan buffer allocation, reducing peak buffer count by 30 to 48%, and device-affinity scheduling, reducing NPU-CPU transitions by 42 to 65%. Across six model families ranging from 125M to 8B parameters, evaluated on WikiText-103 and GLUE, Forge-UGC delivers 6.9 to 9.2x faster compilation than OpenVINO and ONNX Runtime, 18.2 to 35.7% lower inference latency, and 30.2 to 40.9% lower energy per inference. Fidelity is preserved, with max absolute logit differences below 2.1e-5 and KL divergence below 8.4e-9. We also introduce Fusion Gain Ratio, Compilation Efficiency Index, and per-pass execution profiling for systematic evaluation of NPU compilation pipelines.