@JeffDean Wishing you the absolute best for all you are pursuing next, and I'm sure it'll change the world again. I also remain incredibly grateful for hosting me as visiting faculty at Brain in 2018-19 -- that led to MLIR.
Dynamic sequence lengths would lead to two recompilations at most with Inductor (I think); it's a one-time cost that shouldn't impact benchmarking. Programmers would ideally try out torch.compile before resorting to developing/using kernels; the cost of trying out torch.compile is nearly zero if it's known which parts to compile if applied to specific parts. So it'd be valuable/interesting to see the gap b/w torch.compile (Inductor) and kernels here.
Nice thread! For completeness, for the result chart above, do you have the performance bar with torch.compile applied to those same functions that were replaced with kernels? (And also with torch.compile on the whole model ideally with CUDA graphs). Or is "No kernels" already with torch.compile?
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
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
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
CC is co-located with HPCA, PPoPP, and CGO. Great opportunity to be in sync with a large area of computer systems spanning high-performance computer architecture, compiler systems, parallel programming, HPC, and code optimization! https://t.co/QGarszEbbg
CC is co-located with HPCA, PPoPP, and CGO. Great opportunity to be in sync with a large area of computer systems spanning high-performance computer architecture, compiler systems, parallel programming, HPC, and code optimization! https://t.co/QGarszEbbg
The ACM SIGPLAN International Conference on Compiler Construction (CC) 2026 is inviting papers in areas capturing modern compiler design/construction, and its adjacent areas. For more information: https://t.co/eQj0y6lGMh
CC 2026 will be held in Sydney, Australia, 31-Jan, 1-Feb.
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
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
Among all the reasons you listed in your nice article on this, I think a key one is missing: the company in context building the chip isn't itself building an MLIR-based AI compiler for higher-order programming frameworks (cuTile was after nearly 20 years, and it's still only "low/mid-level" abstraction). By the time new hardware is out and easily available, hand-tuned highly-optimized cu* libraries are already available (written in CUDA, CUTLASS, and perhaps inline PTX and assembly as well), including composed library methods for common things in LLMs. So "outside" compilers would be catching up forever! Things would have been very different had NVIDIA or similar chip companies built an MLIR-based AI compiler (for, say, PyTorch, JAX, etc.) **concurrently** with their investments in CUDA-like programming models.
I don't think "fragmentation", lack of upstream reference implementations, or governance are the real or important reasons. There is enough freedom to innovate in downstream MLIR projects. The technical approach to compilation that could surpass what's done by hand is actually an important one, though.
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
@akssri0 Their software execution strategy is pretty poor! See this amazingly detailed article from SemiAnalysis that painstakingly describes all the potholes and provides a detailed list of recommendations to AMD leadership on what to fix and how. https://t.co/taTwYQbN6D
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
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: https://t.co/NeOmZHpbMj
#compiler #ai #polyblocks
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
Great "LLVM in 2024" summary from @michaellarabel at @phoronix: More Than 37k Commits from 2,138 contributors adding up to 35.5 Million Lines of code. Incredible teamwork from an extremely collaborative community of dragons! 🐉
https://t.co/2m6SoOGrkn