Deep-tech startup. Compilers and auto code generators for high-performance AI. PolyBlocks. High-performance computing with polyhedral and MLIR-based compilers.
Bengaluru, India
Joined April 2023
- Tweets81
- Following2
- Followers646
- Likes27
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!
ALT Polyblocks performance with CUDA Tile backend - wall clock times.
ALT Polyblocks performance with CUDA tile backend - kernel times.
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 Labs retweeted
Don't miss coverage of #HiPEAC26 by our media partners, @eetimes in this recap of day 1 👉 eetimes.com/hipeac-2026-dont…
which includes takeaways from Michaela Blott's keynote (such as how to address RAMageddon), along with interviews with Micol Spitale (@polimi) and Uday Reddy Bondhugula (@polymagelabs).
Nitin Dahad (@ndahad) and Pablo Valerio (@PABL0VALERIO) have been getting up to date on the latest advances in computing systems during the event and are posting their articles and videos. We'll be sharing more shortly – watch this space 👀
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…
ALT PolyBlocks compiler stack: from PyTorch/JAX/etc. to hardware.
ALT Automatic fusion with the PolyBlocks compiler for PyTorch to Tenstorrent hardware.
ALT Automatic fusion and tiling for the attention layer with the PolyBlocks compiler for PyTorch to Tenstorrent hardware.
ALT The PolyBlocks team -- at IISc, Bangalore, India.
PolyMage Labs retweeted
Polymage Labs Ties Up with Tenstorrent to Build Compiler for AI Hardware
PolyMage Labs, a deep-tech startup incubated at the Indian Institute of Science (IISc), has entered into a strategic collaboration with US-based Tenstorrent Inc to develop an AI compiler for Tenstorrent's AI accelerators.
The partnership brings together Polymage Labs' compiler expertise and Tenstorrent's AI computing platforms, with the aim of improving software support for advanced AI hardware. Both companies said the collaboration focuses on addressing long-standing challenges in making specialised AI chips easier to program and deploy.
Read more: 2cm.es/1mkAo
#PolymageLabs #Tenstorrent #AIHardware #AICompilers #DeepTech #ArtificialIntelligence #Semiconductors #TechCollaboration #AIInfrastructure #FutureOfAI #MakeInIndia #GlobalTech #StartupIndia
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
ALT PolyBlocks compiler performance on PyTorch AI workloads.
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
ALT PolyBlocks vs Triton vs Cublas/cutlass
PolyMage Labs retweeted
An interesting article that, besides the details, shows why AI chip vendors need to open programming doors at multiple levels:
1) high-level Python-based programming framework support (e.g. PyTorch, JAX),
2) compiler support for (1) that could be turned off/on (e.g. torch.compile),
3) mid/low-level programming support (e.g., CUDA, CUTLASS, Triton),
4) low-level MLIR dialects that expose their hardware intrinsics/virtual ISA on top of which both "(2) compilers" and "(3) low-level frameworks" rest, and
5) the ability to use inline virtual ISA if really needed (e.g. PTX).
While (5) isn't strictly necessary, major chip vendors are missing big time on (4) and hence setting themselves up for failure on (1) through (3). Innovations like those by DeepSeek would have otherwise either not been possible or been incredibly slower/harder!
NVIDIA has excelled at all of these, either doing it themselves or enabling others (open-source community, academia, or the industry) to build it for their hardware. An important message for all other chip vendors who have non-existent or poor support at several of those five levels, or worse, are debating and pitting one programming approach/level against the other! Instead, get the lower-level building blocks right to enable (1) through (5).
medium.com/@amin32846/unlock…
PolyMage Labs retweeted
Automatically generated compiler code at 85% of FlashAttention performance over 45 different attention layer problem sizes/configurations (geomean) drawn from various AI models! About 1.4x as fast in several cases!
All indications are that these improvements will widen and the remaining gaps will close with general and better optimization in the mid and lower levels of MLIR. A sustainable and scalable approach to deliver higher performance!
Automatically generated code from PyTorch-PolyBlocks is now 1.3 to 1.4x as fast as hand-written FlashAttention kernels in several cases! At 85% of Flash-Attention performance overall (geomean) across 45 diverse attention layer problem sizes/configurations drawn from numerous state-of-the-art AI models!
Results below are with PyTorch 2.5 on an NVIDIA A100. (FlashAttention was benchmarked via Torch Inductor.)
More info: docs.polymagelabs.com
#compiler #ai #polyblocks
ALT PolyBlocks (automatic) vs FlashAttention benchmarked on an NVIDIA A100
Automatically generated code from PyTorch-PolyBlocks is now 1.3 to 1.4x as fast as hand-written FlashAttention kernels in several cases! At 85% of Flash-Attention performance overall (geomean) across 45 diverse attention layer problem sizes/configurations drawn from numerous state-of-the-art AI models!
Results below are with PyTorch 2.5 on an NVIDIA A100. (FlashAttention was benchmarked via Torch Inductor.)
More info: docs.polymagelabs.com
#compiler #ai #polyblocks
ALT PolyBlocks (automatic) vs FlashAttention benchmarked on an NVIDIA A100
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
ALT PolyBlocks compiled code: performance improvement with WMMA-level fusion
PolyMage Labs retweeted
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…
ALT PyTorch model coverage with the PolyBlocks compiler backend.
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…
ALT PyTorch model coverage with the PolyBlocks compiler backend.
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
ALT TensorRT vs PolyBlocks with PyTorch on batch size 1
ALT TensorRT vs PolyBlocks with PyTorch on batch size 8
PolyMage Labs retweeted
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
ALT PolyBlocks vs CuDNN on various convolution operator sizes.
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
ALT PolyBlocks vs CuDNN on various convolution operator sizes.
PolyBlocks-compiled PyTorch model on an NVIDIA Jetson Orin Nano (embedded system), executing DETR (transformer-based object detection) on every video frame: 20x as fast as anything else that worked out of the box!
youtube.com/watch?v=ei9zNbyT…
PolyBlocks can now compile out of the box for an edge AI device as well—in this case, an NVIDIA Jetson Orin Nano (Arm64 + integrated NVIDIA GPU). A couple of large HuggingFace models benchmarked here!
- About 2x as fast as Torch 2.4.0 (eager) on DETR (Detection transformer for object detection) and YOLOS.
- Can generate a native binary for ARM64 + integrated GPU in PolyBlocks AOT mode
In action on a video: youtube.com/watch?v=ei9zNbyT…
Acks: Thanks to 'dusty-nv' github.com/dusty-nv for the PyTorch 2.4.0 CUDA 12.2 wheel for Jetson.
ALT Compile Torch with PolyBlocks with a single line.
ALT Torch (eager) vs Torch/PolyBlocks compiled on an NVIDIA Jetson Orin Nano.