PyTorch and torch.compile

From eager Python to optimized GPU kernels

Learning goals

This page will document how PyTorch represents and executes machine-learning programs, and how torch.compile transforms those programs into optimized code. My goal is to connect the high-level model code I write in Python to the graphs, compiler passes, and GPU kernels that execute it.

• Understand eager execution, tensors, autograd, modules, operators, and dispatch.
• Trace the compilation path through TorchDynamo, FX, AOTAutograd, and TorchInductor.
• Inspect the generated graphs and Triton or C++ kernels.
• Recognize graph breaks, guards, recompilations, and dynamic-shape behavior.
• Measure compilation overhead and steady-state training or inference performance correctly.
• Integrate a custom C++ or CUDA operator without breaking autograd or compilation.

Study plan

1. PyTorch execution fundamentals

I will begin with the eager execution model and follow an operation through tensors, autograd, modules, the operator library, and device dispatch. This provides the baseline for understanding what changes when compilation is enabled.

2. The torch.compile stack

• TorchDynamo: capture Python execution into FX graphs.
• AOTAutograd: capture and partition forward and backward graphs.
• TorchInductor: schedule operations, fuse work, and generate optimized code.
• Triton: generate and execute many of the GPU kernels emitted by Inductor.

3. The compiler programming model

I will study graph breaks, guards, cache behavior, recompilation, symbolic and dynamic shapes, compilation modes, and the boundaries between compiled and eager execution.

4. Debugging and performance analysis

I will use compiler logs and the PyTorch profiler to inspect captured regions, generated code, kernel launches, graph breaks, and CPU or GPU gaps. Benchmarks will separate initial compilation cost from warm steady-state execution.

5. Extending the stack

The final stage will connect the compiler track to my CUDA work by registering and testing a custom operator, integrating it into a compiled model, and comparing it against an eager baseline and Inductor-generated kernels.

Experiments

• First experiment: compare eager and compiled execution for a small model, including warm-up, compilation time, and steady-state latency.
• Introduce a graph break deliberately, locate it with compiler logs, and rewrite the code to remove it.
• Inspect an FX graph and the corresponding Inductor-generated Triton kernel.
• Test changing input shapes and explain when guards cause recompilation.
• Register a custom C++ or CUDA operator and verify correctness with eager execution, autograd, and torch.compile.

Notes and results

Notes, profiler traces, generated code, benchmarks, and conclusions will be added here as I complete each experiment.

Resources

• Introduction to torch.compile.
• torch.compiler overview and compiler-stack components.
• torch.compile core concepts and programming model.
• torch.compile troubleshooting and compiler logs.
• Profiling torch.compile performance.
• PyTorch custom operators.
• GPU MODE lecture materials.