PolyMage Labs

Deep-tech startup. Compilers and auto code generators for high-performance AI. PolyBlocks. High-performance computing with polyhedral and MLIR-based compilers.

[PolyMage Labs @polymagelabs]

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.
https://t.co/PmJ9RTA8BI

[HiPEAC @hipeac]

Don't miss coverage of #HiPEAC26 by our media partners, @eetimes in this recap of day 1 👉 https://t.co/bhjZg7gePt
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 Labs @polymagelabs]

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: https://t.co/ncvAP1GO12

[Entrepreneur Asia Pacific @EntrepreneurAPJ]

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: https://t.co/Op5nYuW9tx

#PolymageLabs #Tenstorrent #AIHardware #AICompilers #DeepTech #ArtificialIntelligence #Semiconductors #TechCollaboration #AIInfrastructure #FutureOfAI #MakeInIndia #GlobalTech #StartupIndia

[PolyMage Labs @polymagelabs]

PolyBlocks is not yet ready to be open source, but please see below on how to use/collaborate with it, especially if you are (1) a hardware vendor who wants a PolyBlocks backend/target built for their AI chip (highly valuable!), or (2) someone who wants to generate fast code/kernels for existing NVIDIA/AMD GPUs from PyTorch/JAX functions.
https://t.co/1FwftGLTVB

[PolyMage Labs @polymagelabs]

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: https://t.co/NeOmZHpbMj

[PolyMage Labs @polymagelabs]

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 https://t.co/LKRWAmgsV4

Results here show high performance can be achieved with compact/productive programming if compilers are built right!

Know more: https://t.co/NeOmZHpbMj

[Uday Bondhugula @ukrb84]

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).
https://t.co/rkf9W92kP4

[Uday Bondhugula @ukrb84]

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!

[PolyMage Labs @polymagelabs]

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: https://t.co/NeOmZHpbMj
#compiler #ai #polyblocks

[PolyMage Labs @polymagelabs]

Nsight compute statistics comparing the instruction profiles notably show the difference between LDS and STS instructions! Note the increased utilization of the math pipeline after elimination of shared memory transfers!

[PolyMage Labs @polymagelabs]

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: https://t.co/NeOmZHpbMj

[PolyMage Labs @polymagelabs]

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 (https://t.co/D8aYnJgPrh). 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: https://t.co/EXsKSpMgsU

[PolyMage Labs @polymagelabs]

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 (https://t.co/D8aYnJgPrh). 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: https://t.co/EXsKSpMgsU

[PolyMage Labs @polymagelabs]

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!
https://t.co/k34xDdr9I1
Docs: https://t.co/NeOmZHpbMj

[PolyMage Labs @polymagelabs]

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: https://t.co/NeOmZHpbMj

[PolyMage Labs @polymagelabs]

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: https://t.co/NeOmZHpbMj

[PolyMage Labs @polymagelabs]

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!
https://t.co/iSaaRJUhyh

[PolyMage Labs @polymagelabs]

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: https://t.co/6qRV7CXY5L

Acks: Thanks to 'dusty-nv' https://t.co/HGQK3ioR6Q for the PyTorch 2.4.0 CUDA 12.2 wheel for Jetson.