How Pytorch 2.0 accelerates deep learning with operator fusion and CPU/GPU code-generation

How Pytorch 2.0 Accelerates Deep Learning with Operator Fusion and CPU/GPU Code-Generation

A primer on deep learning compiler technologies in PyTorch for graph capture, intermediate representations, operator fusion, and optimized C++ and GPU code generation

Computer programming is magical. We write code in human readable languages, and as though by magic, it gets translated into electric currents through silicon transistors making them behave like switches and allowing them to implement complex logic — just so we can enjoy cat videos on the internet. Between the programming language and hardware processors that run it, is an important piece of technology — the compiler. A compiler’s job is to translate and simplify our human readable language code into instructions that a processor understands.

Compilers play a very important role in deep learning to improve training and inference performance, improve energy efficiency, and target diverse AI accelerator hardware. In this blog post I’m going to discuss deep learning compiler technologies that powers PyTorch 2.0. I’ll walk you through the different phases of the compilation process and discuss various underlying technologies with code examples and visualizations.

What is a deep learning compiler?

A deep learning compiler translates high-level code written in deep learning frameworks into optimized lower level hardware specific code to accelerate training and inference. It finds opportunities in deep learning models to optimize for performance by performing layer and operator fusion, better memory planning, and generating target specific optimized fused kernels to reduce function call overhead.

Unlike traditional software compilers, deep learning compilers have to work with highly-parallelizable code often accelerated on specialized AI accelerator hardware (GPUs, TPUs, AWS Trainium/Inferentia, Intel Habana Gaudi etc.). To improve performance, a deep learning compiler has to take advantage of hardware specific features such as mixed precision support, performance optimized kernels and minimize communication between host (CPU) and AI accelerator.

While deep learning algorithms are continuing to advance at a rapid pace, hardware AI accelerators have also been evolving alongside to keep up with deep learning algorithm performance and efficiency needs. I discuss the co-evolution of algorithms and AI accelerators in an earlier blog post:

and articles have of different models from huggingface, timm and torchbench.

At a high-level the default options for PyTorch 2.0 deep learning compiler performs the following key tasks:

1. Graph capture: Computational graph representation for your models and functions. PyTorch technologies: TorchDynamo, Torch FX, FX IR
2. Automatic differentiation: Backward graph tracing using automatic differentiation and lowering to primitives operators. PyTorch technologies: AOTAutograd, Aten IR
3. Optimizations: Forward and backward graph-level optimizations and operator fusion. PyTorch technologies: TorchInductor (default) or other compilers
4. Code generation: Generating hardware specific C++/GPU Code. PyTorch technologies: TorchInductor, OpenAI Triton (default) other compilers

Through these steps, the compiler transforms your code and generates intermediate representations (IRs) that are progressively “lowered”. Lowering is a term in the compiler lexicon that refers to mapping a broad set of operations (such as supported by PyTorch API) to a narrow set of operations (such as supported by hardware) through automatic transformation and re-writing by the compiler. The PyTorch 2.0 compiler flow:

If you are new to compiler terminology don’t let all of this scare you yet. I’m not a compiler engineer either. Keep reading and things will become clear as I’ll break the process down using a simple example and visualizations.

A walk through the torch.compile() compiler process

Note: This whole walkthrough is in a

The output of that above code snippet are the FX IR code and the graph diagram showing our function sin^2(x)+cos^2(x)

Note that our fake compiler inspect_backend function is only invoked when we call the compiled function with some data i.e. when we call compiled_model(x). In the above code snippet, we’re only evaluating the function or in deep learning terminology, doing a “forward-pass”. In the next section we’ll take advantage of the PyTorch’s automatic differentiation engine called torch.autograd to compute the derivative and the “backward-pass” graph.

Automatic differentiation: Forward and backward computational graphs

PyTorch technologies: AOTAutograd, Core Aten IR

TorchDynamo gave us the forward pass function evaluation as an FX graph, but what about the backward pass? For the sake of completeness, I’m going to digress from our primary topic and talk a bit about why we need to evaluate the gradients of a function with respect to its weights. If you’re already familiar with how mathematical optimization works skip this immediate section.

What is backward pass and backward graph?

The “learning” part of deep learning and machine learning is a mathematical optimization problem which is simply stated as: Find the value of a variable w that yields the lowest value of some function of w. Or more succinctly:

In machine learning f(w) is the loss function parametrized by weights. f(w) can be more clearly represented as some measure of error between the training labels and the model’s prediction labels based on the training data:

Turns out, if we can calculate the “rate of reduction” of loss with respect to weights, we can update our weights to move one step closer to a smaller and smaller loss f(w). In other words, we must move closer to a model that better fits our training dataset. We can find next values of weights by calculating the steepest slope of the loss f(w) at a given w and perturb w to head in that direction. The slope of a function with respect to the weights, is its derivative with respect to the weights. Since there are more than one weight values, the derivative becomes a vector quantity called the gradient which is a vector of partial derivatives with components for each weight. The weights w are perturbed at each iteration by some function g() of the gradients as follows:

Where the function g(.) depends on the optimizer (e.g. sgd, sgdm, rmsprop, adam etc.).

For SGD the weight update step becomes:

How does PyTorch 2.0 trace the backward pass graph?

First let’s calculate what we expect the backward pass graph should look like and then compare it with what PyTorch generates. For our simple function, the forward graph and the backward graph should implement the following function. If sin and cos bother you, you can imagine f(x) being the loss function applied to a neural network.

PyTorch uses reverse-mode

This will generate the following forward AND backward graphs. Notice that the forward graph also looks slightly different from what we saw earlier in Figure x. For example torch.sin(x) in the FX graph IR and in our original code has been replaced by torch.ops.aten.sin.default(). What’s this funny thing called aten, you might ask, if you’re not already familiar with it. ATen stands for A Tensor library, which is a very creatively named low level library with a C++ interface that implements many of the fundamental operations that run on CPU and GPU.

In eager mode operation, your PyTorch operations are routed to this library which then calls the appropriate CPU or GPU implementation. AOTAutograd automatically generates code that replaces the higher level PyTorch API with ATen IR for the forward and backward graph which you can see in the output below:

You can also see that in addition to the output of the forward pass, the forward graph outputs some additional tensors [add, sin, cos, primals_1] . These tensors are saved for the backward pass for gradient calculation. You can also see this in the computational graphs for the forward and backward pass in the figure shared earlier.

What are the different types of IR in PyTorch?

ATen IR is a list of operators supported by the ATen library as we discussed in the previous section, and you can see the full list of operations implemented in (formerly canonical Aten IR) is a subset of the Aten IR that can be used to compose all other operators in the Aten IR. Compilers that target specific hardware accelerators can focus on supporting only the Core Aten IR and mapping it to their low level hardware API. This makes it easier to add hardware support to PyTorch since they don’t have to implement support for the full PyTorch API which will continue to grow with more and more abstractions.

and

Notice that I’ve passed optional argument that enables two debug features:

• trace.enabled: Generates intermediate code to inspect code generated by TorchInductor
• trace.graph_enabled: Generates the optimized computational graph visualization after operator fusion

For our simple example TorchInductor is able to fuse all intermediate operations in our function into a single custom operator, and you can see below how that simplifies the forward and backward computational graphs.

You must surely be wondering what this fused operator looks like in code. The code for the fused operator is automatically generated by TorchInductor and it’s in C++ or Triton based on the target device — CPU or GPU. You don’t need to explicitly specify to TorchInductor which device to target, it can infer it from the data and model device type.

To view the generated code, you have to enable debugging using trace.enabled=True and this creates a director called torch_compile_debug with debug information.

The full path to forward and backward graph code are:

• torch_compile_debug/run_<DATE_TIME_PID>/aot_torchinductor/model__XX_forward_XX/output_code.py
• torch_compile_debug/run_<DATE_TIME_PID>/aot_torchinductor/model__XX_backward_XX/output_code.py

If you set device = ‘cuda’ (assuming your computer has a GPU device) then the generated code in the forward folder is in Open AI Triton

If you set device = ‘CPU’ then the generated code is in C++ with OpenMP pragmas

Recap

Deep learning compilers are complex with intricate inner workings that rival a swiss watch. In this blog post, I hope I provided you with a gentle and easy to follow primer on this topic and how these technologies power PyTorch 2.0.

1. I started with a simple PyTorch function
2. I showed how TorchDynamo captures the graph and represents it in FX IR
3. I showed how AOTAutograd generates the backward pass graph, lowers PyTorch operators into Aten operators and represents it in an FX graph container.
4. I discussed how Aten operators can be further decomposed into Core Aten IR and Prims IR that reduces the number of operators that other compilers can support without supporting the full PyTorch API list.
5. I showed how TorchInductor performs operator fusion and generates optimized code for CPU and GPU targets

If you followed along you should be able to provide a high-level response to the following questions:

• What is a deep learning compiler?
• What does a PyTorch 2.0 compiler do when you call torch.compile()?
• Why do we need a forward and backward pass graph ahead of time?
• What are the different intermediate representations (IR) in PyTorch
• What is the difference between ATen IR, Core ATen IR, Prims IR?
• What is operator fusion and why is it important?

Thank you for reading all the way to the end!

If you found this article interesting, consider following me on medium to be notified when I publish new articles. Please also check out my other blog posts on ), was originally published in Towards Data Science on Medium, where people are continuing the conversation by highlighting and responding to this story.