When you write a program (say, in C or C++), it can’t be interpreted by a computer in its raw form - before your CPU can execute your program, it needs to compile your code with a compiler to translate your human ‘source code’ into computer-readable ‘machine code’.
However, PyTorch (which is built on top of Python and runs on your GPU) doesn’t actually need compilation to run for its specific operations. By default, PyTorch uses eager execution - wherein operations are immediately dealt with as they occur going through the program sequentially, rather than the entire program being compiled as a whole before being run. When dealing with a specific operation, PyTorch launches a precompiled kernel which actually does the work on your GPU.
Many operations in modern deep learning are going to launch multiple separate kernels, and thus this method of launching kernels can be extremely inefficient:
- For one, launching an individual kernel has its own fixed overhead. Then, logically, we want to limit the launching of new kernels as much as possible.
- Secondly, in many scenarios, a lot of your total latency is running back and forth from your GPU’s compute units to your memory transferring data. If you’re working with the same data in multiple operations, then writing data to your GPU’s memory unnecessarily slows your program down.
This is really simplified, but take:
x = a + b
y = x * 2
X is being used in both operations, so we could save ourselves a trip down memory lane (get it?) by just running both operations in a single kernel.
Modern GPUs use a memory hierarchy, where there are multiple levels of memory (individual thread registers -> block chip shared memory -> L1 cache (for your Streaming Multiprocessor) -> L2 cache across all SMs -> Global memory -> technically all the way to your CPU memory).
These are labeled in sequential time to access - going back and forth to global memory / GMEM is EXTREMELY slow! We must manage our memory properly.
The PyTorch compiler, at a high level, captures some sequence of operations, and seeks to optimize your program, generating more specialized and fast code that is better at executing.
The ‘lifecycle’ of a PyTorch program goes (roughly) as follows:
TorchDynamo
Like a more traditional compiler, the PyTorch compiler still needs to take your program and translate it into some format in which it can optimize things. For compilers, this is known as an Intermediate Representation or IR.
Compiler optimization can be dumbed down to trying to transform your IR to however you so please (getting rid of dead code, optimizing memory access, etc.).
As a side note, one IR might not be blanket-applicable for most scenarios; one modern solution is MLIR (Multi-Level Intermediate Representation), a framework part of the LLVM compiler that tries to maintain the right level of abstraction at the right time.
TorchDynamo uses an IR called an FX graph, a computation graph of sorts that explains the relationship between the operations and how data moves. It generates this by intercepting Python’s regular execution and symbolically tracing the program. Rather than eagerly evaluating the expressions, it records what operations are happening and their dependencies, constructing this into a DAG (directed acyclic graph).
Whenever this symbolic evaluation engine can’t understand some instruction (how dare it), it introduces a graph break, which essentially hands that portion of code to the interpreter (executing it like normal) before resuming trying to create this graph after the break.
AOTAutograd
Not to be confused with Autograd!
Once we have an FX graph from TorchDynamo, the next step is AOTAutograd, which prepares the FX graph for the backend of the compiler.
PyTorch’s similarly named Autograd (the automatic differentiation engine) builds a graph of gradient dependencies and uses this to calculate gradients during the backwards pass / backprop.
AOTAutograd (Ahead of Time Autograd) is different - essentially, your FX graph is too user-facing to later become kernels. Before we hand this graph off, we need to do some normalization, which is the role of AOTAutograd. Normalization is important because PyTorch has lots of operations, many of which are pretty complicated and hard to easily optimize.
First, operations that overwrite existing memory are removed. If some tensor is modified throughout the runtime of a program, the compiler needs to keep track of which time it is to keep track of the correct value of the tensor. As a result, tensors become functionally immutable, in the sense that their values are unchanged after being created. Instead, these operations are represented as new tensor values.
Next, we try to reduce the complexity in the total span of operations possible. Each complex operation can be represented as some ‘linear combination’ of sorts of more fundamental operations. That’s what we do here - we represent all operations as some combination / permutation of smaller, more standardized operations called ATen operators. Some ATen operators are add (addition), mul (multiplication), sum (self-explanatory), and so on.
I almost think of this as analogous to how a C/C++ program is translated to Assembly instructions. The fall of my freshman year, I had to hand-translate a C program into ASM for a class project (CMSC216 Computer Systems); all it is is breaking up your bigger functions into smaller ones, basically what the computer is doing here!
If you’re training a network, AOTAutograd also is responsible for generating a graph of the backwards pass (of which it uses the existing automatic differentiation system to do so). This allows for the backward pass to be compiled/represented/optimized in the same means as the forward pass is.
Now that we have our graph, we pass it to TorchInductor!
TorchInductor
TorchInductor essentially turns your graph into some set of very fast very performant kernels. There are multiple steps to this:
Lowering
We already ‘lowered’ our program into an intermediate representation (basically the main point of the prior stages). So why lower again? Well, if we can split a single value into multiple, like so:
# old representation - a and b are tensors that can be iterated elementwise!
x = a + b
# new representation
for i in range(N):
x[i] = a[i] + b[i]
Now, we both have a more concrete and well-defined version of the computation and a means of knowing how individual subelements are calculated. This makes it very easy to perform multiple of these at once in parallel - what our GPU is great at!
Scheduling
This is the most interesting part of the torch compiler to me - I plan to write more about this component specifically in the future, so stay tuned! Essentially, the scheduler decides how to organize our massive set of computations into kernels (functions you can run in parallel on your GPU).
Let’s return to our very first code example from the start:
# think about this as a kernel or individual function
# old representation
x = a + b
y = x * 2
# representation in TorchInductor's IR - Loop IR
for i in range(N):
x = a[i] + b[i]
for i in range(N):
y[i] = x[i] * 2
# NEW representation - fusing these into one kernel!
for i in range(N):
y[i] = (a[i] + b[i]) * 2
Of course, this is overly simplistic, but you can see how we were able to avoid the intermediary read/write of the X tensor to GPU memory by fusing these kernels into one. This saves us time. Logically, if fusing kernels saves us time to avoid memory read/writes, can’t we just fuse everything into a single kernel?
For some cases: technically you can! Recent work on Megakernels does this - some wonderful papers / blogs / work from Flapping Airplanes / Stanford’s Hazy Research / Cohere has been published recently about these. We’ve been able to fit entire forward passes or decoding steps in single kernels - there’s definitely more, but this list is only off the top of my head. I also want to learn more / blog about a Megakernel soon, more on that to come!
But anyways, can we just fuse everything into a megakernel or some other large kernel? Not really. For one, not all fusion is possible: kernels that can be productively fused tend to share traits (for example, accessing the same data or doing the same operations). For kernels A and B, their execution requirements must allow them to run together in the first place.
Additionally, not all fusion is beneficial. Overfusing can stress thread registers, do duplicate work, or limit the amount of computation that can be done in parallel.
It’s also to note that fusing kernels is not always beneficial, it’s mainly beneficial in cases of being memory-bound and not compute-bound (although with compute-bound scenarios it reduces overhead).
To formalize, TorchInductor uses can_fuse() to see if some fusion is possible, and of the candidates that can, it uses score_fusion() to rank all of the fusion choices. From there, it greedily fuses the top ranking candidates.
Some heuristics it uses include:
- Amount of memory traffic saved (bytes)
- Graph proximity on computation graph
- Operation type
- Shared data
- End size of fused kernel (bytes)
It returns some FusionScore object which is comprised of several of these measurements, and then greedily fuses the ‘best’ combinations. Fusing kernels together changes the ‘topology’ of the landscape of possible fusions - that is to say it changes the possibilities of kernels that you can fuse. This process is run iteratively in up to 10 rounds.
CodeGen
Once we have decided what to fuse, we turn these fusion choices into executable kernels in Triton (a language developed to write GPU kernels by OpenAI; if CUDA = C++, Triton = Python so to speak, it’s slightly more abstracted away):
@triton.jit
# Knowing triton syntax is not necessary -
# but this is what our resulting kernel could formally look like:
def kernel(A, B, Y, N, BLOCK: tl.constexpr):
i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
mask = i < N
a = tl.load(A + i, mask=mask)
b = tl.load(B + i, mask=mask)
result = (a + b) * 2
tl.store(Y + i, result, mask=mask)
The specific things like tile shapes (deciding how tiles map onto the compute, and other things) are decided by Inductor. Triton has its own compiler, which it uses to compile the kernels (obviously), where the compiler chooses things like the mapping of tensor elements onto GPU threads and how instructions use memory.
Afterwards, there’s ‘machine code’ for your GPU to execute.
It’s also important to note that Inductor does not automatically create Triton kernels for everything. The engineers at NVIDIA have always been very smart, and human hyper-optimized libraries for specific tasks like GEMMs (General Matrix Multiplications) have already existed. For these tasks where some optimized implementation exists already, the compiler simply reuses it.
Runtime
A generated wrapper launches the kernels that were made. NVIDIA GPUs contain individual Streaming Microprocessors (SMs), which manage thread blocks (groups of individual threads). So when your kernel launches a grid of thread blocks, the GPU scheduler decides which blocks to assign to certain SMs.
Kernel compilation / other metadata is cached to not have to recompute for a later call.
Conclusion
Technically, torch.compile() does induce extra computational overhead, although it does generally slash execution time.
But that’s a high-level overview of how the PyTorch Compiler / torch.compile() works. Torch compiler is different than other ML compiler frameworks for other companies: Google, for example, developed its XLA (Accelerated Linear Algebra) framework which does graph-based things, unlike the Just-In-Time / JIT compilers like Inductor.
I also did mention the MLIR frameworks like with LLVM’s framework. But the point is that the compiler for PyTorch is really interesting (and really optimized).
This was a fairly high-level overview where I glossed over lots of important implementation details - I plan to explore TorchInductor’s scheduling / fusion heuristics in more detail in a later post!