Uday Bondhugula

AI for principled compiler construction! Everyone is talking about AI-generated high-performance kernels, but the real opportunity lies in using AI effectively to build compilers! Why use AI to generate kernels one by one when it can build a kernel generator for infinite reuse? By pairing solid compiler engineering principles with tools like Claude, we built a CUDA Tile IR backend for PolyBlocks in a few weeks, achieving a level of performance that typically takes multiple human-years! Our backend now compiles unmodified PyTorch and JAX models to a sequence of fused CUDA tile kernels (all via MLIR) with state-of-the-art performance. By layering on top of the NVIDIA CUDA Tile IR compiler, we offload lower-level optimizations to it. Both NVIDIA GPU backends of PolyBlocks, cuda-tile-based and non-cuda-tile-based (which we had developed over several years), still rest on the same Polyblocks code generation infrastructure for affine fusion, tiling, mapping to MMAs, 2-d tilization, etc., but the cost models they use on both paths are slightly different. Check out the benchmarks below showing wall-clock and total kernel execution times. More to come! Compilers are now finally evolving rapidly!

Finally, here's the paper on PolyBlocks describing how fully code-generating compilers for AI chips can be built! This is the culmination of multiple years of R&D and engineering. There is now enough reusable infrastructure in our toolkit to quickly build high-performing PyTorch/JAX compilers for new chips, no matter how weird or unique their capabilities are, and without relying on any "kernel" libraries or manual model optimization or porting. The paper isn't exhaustive, but it provides details on the key parts, the design choices, and why they are powerful. arxiv.org/abs/2603.06731

Polymage's PolyBlocks engine enables the rapid development of compilers for new AI chips, extracting close-to-peak performance from hardware, and realizing high-performance AI systems for vision, text, generative AI, and more. Highly promising results here from Polymage's collaboration with @tenstorrent --- showing that the right compiler technology can be a leapfrog advance to the programming ecosystem for new chips! More details: docs.polymagelabs.com/articl…

With constant improvements in its optimization passes, PolyBlocks-generated code continues to get better! We compare here against all other state-of-the-art production PyTorch backends on a variety of workloads, including transformer-based and convolution-based ones. The latest round of improvements comes from more/better unroll-and-jam, more fusion across operators, simplifying away invariant computation, and wider use of on-chip memory. The impact of such compiler technology would be even greater for new hardware that doesn't yet have a well-established library/programming ecosystem. More info: docs.polymagelabs.com

Programming and compiling for AI hardware happen today through multiple approaches, but they broadly fall into low-level, mid-level, and high-level ones. (1) Low-level is based on hand-written kernels (CUDA, CUTLASS, assembly/PTX, similar) on which the default execution engines of AI frameworks (eager-mode) are built; (2) mid-level is mostly OpenAI Triton, Pallas, and cuTile (announced); and (3) high-level ones, which are compilers for PyTorch/JAX/TF functions, e.g. Inductor, XLA, and PolyBlocks. One would stop building low-level libraries/kernels if mid-level ones (e.g. Triton) delivered good/comparable performance; similarly, one would stop relying on hand-written mid-level if high-level compilers delivered comparable performance. Results here compare the high-level fully automatic approach of PyTorch-PolyBlocks against OpenAI Triton and against lower-level kernels (CuBLAS) on an NVIDIA GPU. While Inductor is an automatic compiler, it still relies on a combination of low-level (e.g. cublas/cutlass kernels) and mid-level frameworks (Triton); for matmul, perhaps exclusively cublas/cutlass (so no code generation). For Triton, the implementation used was from triton-lang.org/main/getting… Results here show high performance can be achieved with compact/productive programming if compilers are built right! Know more: docs.polymagelabs.com

Optimizing computations for memory hierarchies (caches, explicitly addressed on-chip memory, and registers) has been an art since superscalar processors became popular in the 90s. Some of the most important computations today are optimized for multiple levels of cache and fast memories all the way up to register files. We show how important the last step of optimizing for register reuse is even after you've optimized everything greatly for on-chip memory, in this case, NVIDIA GPUs' shared memory scratchpads. Results below show the amount of improvement gained by eliminating back-and-forth transfers to shared memory by reusing data in the registers (for all but for the true input and output). This evaluation is done for "attention" layer computations drawn from various AI workloads; these computations are at the heart of all transformer-based networks, LLMs and auto-regressive models in general. The comparison here is between already optimized/fused attention layer code generated by the PyTorch/PolyBlocks compiler with and without WMMA register-level fusion on an NVIDIA A100 - something we can do by turning on/off an optimization flag. On an absolute scale, PolyBlocks compiler-generated code here is at ~80% of hand-written flash-attention-kernel-based performance (measured via Torch inductor), geomean across 26 diverse/assorted attention workloads). More information: docs.polymagelabs.com

While we've posted in the past on PolyBlocks performance, we now report its coverage on all the 210 PyTorch models recently used in the evaluation of the PyTorch-2 paper from ASPLOS 2024 (pytorch.org/assets/pytorch2-…). These benchmarks are publicly available and drawn from HuggingFace, TIMM, and TorchBench suites. PolyBlocks can successfully compile and execute 185 of these 210 models (push button, no source changes)! These models span various AI domains with some containing nearly 25,000 tensor operators. 90% of the models that run with PyTorch's standard runtime or compiler work with PolyBlocks! Benchmarks courtesy: github.com/pytorch/pytorch/t…

For PyTorch, PolyBlocks is now consistently and significantly faster than NVIDIA TensorRT on various AI vision workloads! The PolyBlocks compiler engine recently switched from solely relying on WMMA operations to MMA (these determine how the matmul cores are programmed on NVIDIA GPUs). With this and other improvements, the PyTorch/PolyBlocks compiler consistently beats TensorRT (a state-of-the-art inference engine/compiler employing expert-written kernels) on a variety of deep learning vision models from TorchVision and HuggingFace - by an average of 1.5x and up to 2.3x as fast. Benefits come from better fusion of operators and better-generated code for the convolution operators in several cases. PolyBlocks can be simply used with: torch.compile(func, backend="polyblocks") Try it out on the Polymage playground! playground.polymagelabs.com Docs: docs.polymagelabs.com

The PolyBlocks compiler generates all its code through MLIR, LLVM, and PTX, and we compare it here against the NVIDIA CuDNN library on ~200 assorted deep learning convolution configurations drawn from various popular models (mainly vision). CuDNN is expected to be hand-written by NVIDIA experts in GPU assembly and CUDA. PolyBlocks compiled code is: 1. faster than CuDNN on 73/193 convolutions. 2. provides more than 10% improvement on 43/193 convolutions. 3. within 25% of CuDNN on 147/193 convolutions. 4. within 2x of CuDNN performance on 186/193 convolutions. The performance of automatically generated code benefits from continuous improvements, avoiding specific library maintenance issues arising from the cross-product of multiple versions of kernels tuned for buckets of problem sizes x hardware. More info: docs.polymagelabs.com