PyTorch 2.0: TorchInductor

Key Takeaways

[Music] [Music] thank you hello everyone my name is Jason Ansel and I'm going to talk about torch inductor which is a new compiler back end for pi torch 2.0 when designing torch inductor we started off with three key principles so the first principle is that torch inductor is pi torch native and what this means is that inductor uses very similar abstractions to uh Pi torch eager which allows it to Faithfully capture all the behavior of pytorch the next key principle is python first so with pytorch 2 we are embracing Python and part of that involves writing inductor entirely in Python which makes it much easier to hack on and extend and then finally uh we wanted inductor to be really General and so we focused on breadth rather than depth early on which means tackling the tricky operators and tricky optimizations early to ensure that our design was General and could scale inductor uses uh three key Technologies as well the first technology is a defined by run Loop level AR and what that means is that the core compiler IR that's done at the Loop level is actually a python callable and to do things like code generation and Analysis we actually execute the IR and do that which is a sort of Novel trick in the compiler space but some things that's been used um for pytorch programs uh for a while uh next we we really wanted inductor to support Dynamic shapes and strides um from uh day one and so inductor uses Senpai which is a symbolic math library to reason about uh shapes and generate code that's not uh specialized uh to specific input sizes and finally I'm going to reinvent the wheel we wanted to reuse uh state-of-the-art languages and so we took inspiration from what our users were doing and increasingly we were seeing people write high performance kernels in this new language Triton um by uh Philip tillett at open Ai and uh so what we have is we have a compiler in torch inductor which generates uh Triton code which is easy to understand you can look at the output and and inspect what it does or even change it and on CPUs we generate C plus plus code uh so what has Triton uh Triton is you could think of it as a better Cuda so uh it's a higher level language than Cuda but also lower level than a lot of pre-existing dsls and it makes it really easy to write high performance code even for tricky things like Matrix multiplies that are competitive uh with things like kudi and n and kublas next I'm going to talk about Mine by run Loop level AR and so here's what the IR for a permute and an ad would look like and so we have this inner function here which takes a list of senpai Expressions that list of symbi Expressions represents a coordinate um uh symbolically that we want to generate and then the in the the body of this function we'll call ops.load twice and ops.add once and the way you use this IR is for example if you wanted to write an analysis pass you could replace This Global variable Ops with something which which records the loads then you could run this function and you could uh do your analysis to figure out uh what loads this function makes and if you wanted to do a code generation you could replace this Ops function with something that printed out Trident code execute this this function and do code gen and this makes it really easy to write lowerings because if you're building composite Ops that sort of take two inputs together you can actually use the features of the Python language to build your lowering and then part way through the compilation process we actually use FX to trace this defined by run IR which gives us a IR that's that's more manipulable so here's an overview of the compiler stack uh so uh we start off with the graph that George Dynamo has captured and we use aot autograd and Prim torch to decompose the operator set into a much more minimal operator set around 250 operators and aot autograd captures the forwards and the backwards graph which we compile uh independently next there's uh graph lowerings graph lowerings takes this primitive operator set which is around 250 operators and lowers it to inductors Loop level ir and the inductors Looper level contains only around 50 operators so we've removed a lot of the complexity and simplified the graph substantially uh and uh we're now operating on how to on basically single elements of tensors rather than full tensors uh as in the input next there's the scheduling phase scheduling is where we decide what gets fused with what and do other optimizations such as memory planning and tiling and then finally we have code generation and so there's two pieces to code generation we have the backend code generation which either generates Trident code or C plus plus and then we have what we call wrapper code gen and wrapper Cogen is the code that stitches together the calls to many different kernels and it basically replaces The Interpreter part of the compiler stack so here's some results on GPU I know Sue Miss shared these earlier but inductor generates up to 1.86 Geo mean speed up on large realistic Benchmark Suites and we're super excited about these results and hope to continue to improve them in the future we also have results on on CPU and these CPU results are a collaboration with the Intel Pi torch team we see up to 1.26 X geomine speed up on uh on on CPU inference and this serves the purpose of uh having inductor support CPUs but it also supports the purpose of making sure the inductor is General we didn't want to build a compiler that only supported gpus we wanted something that could scale to support a wide variety of Hardware back ends and having a c plus as well as it is Trident forces that generality so thanks so much uh if you uh the the code base for inductor is in the pi torch repo I've linked it here and there's more information uh in the blog post as well thanks so much [Music]