Post

Torch.compile 101 - From Python Function to Triton Kernel

An interactive deep dive into PyTorch's torch.compile system, tracing the journey from Python functions to optimized Triton kernels with hands-on examples and performance analysis.

Posted Updated
By
20 min read
Torch.compile 101 - From Python Function to Triton Kernel

Badge

Introduction

torch.compile in pytorch 2 is a common and revolutionary approach to optimizing PyTorch eager training performance. As AI Infra engineer, it is common case that we need to deal with benchmarking with different compiling configuration. Dirty tuning is usually fine and sometimes necessary, but more importantly, we want to deep dive and understand the underlying mechanism (which is also critical skills we want to build as system engineer!)

In this interactive deep dive, we’ll trace the journey of a simple Python function through the entire compilation pipeline, from the initial function definition to the final execution of optimized Triton kernels.

🎯 Interactive Goal: By the end of this guide, you’ll be able to run commands and see the actual compilation process in action!

📚 Recommended Reading: This will be a dirty code walkthrough guide. If you want to understand this from theory point of view, the original pytorch 2 paper is the best material in my opinion. Actually, I am using this paper personally as a starter point to understand the overall picture!

Phase 1: TorchDynamo - Graph Capture

What is TorchDynamo?

TorchDynamo is PyTorch’s graph capture mechanism that converts Python functions into computational graphs. It is hacking into the CPython’s compilation stage (reference) and modify the bytecode to be actually executed.

TorchDynamo Architecture

Interactive Example: Let’s Trace a Function

Create a file called trace_example.py:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

# torch version: 2.8.0
import torch
import os

torch.compiler.reset()

def complex_function(x, y):
    """A function that will trigger various optimizations"""
    # Multiple operations that can be fused
    a = x + y
    b = a * 2.0
    c = torch.relu(b)
    e = torch.nn.functional.silu(c) # this will trigger decomposition
    
    # Reduction operations that trigger Triton code generation
    d = e.sum(dim=1)  # Reduce along dimension 1
    f = d.mean()       # Global mean
    
    return f

def simple_function(x, y):
    return x ** 2 + y ** 2

# Create larger tensors to see more optimization
input1 = torch.randn(1000, 1000).to(device="cuda:0")
input1.requires_grad = True
input2 = torch.randn(1000, 1000).to(device="cuda:0")

compiled_fn = torch.compile(complex_function, backend="inductor")

result = compiled_fn(input1, input2)
# result.backward()
print(f"Result: {result}")

Run this example:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

echo "Setting up environment variables for PyTorch Compile debugging..."
export TORCH_COMPILE_DEBUG=1
export TORCH_LOGS=+dynamo,+inductor,+aot,+guards,+recompiles

echo "Environment variables set:"
echo "TORCH_COMPILE_DEBUG=$TORCH_COMPILE_DEBUG"
echo "TORCH_LOGS=$TORCH_LOGS"

echo ""
echo "Running PyTorch Compile Deep Dive Demonstration..."
echo "================================================"

rm -rf /tmp/torchinductor_jobuser # cleanup the torchinductor cache

python trace_example.py

echo ""
echo "Demo completed! Check the generated debug files in torch_compile_debug/"

What you’ll see (under torch_compile_debug/run_xxx/torchdynamo/debug.log):

  • Detailed bytecode tracing output showing each instruction
  • FX graph construction with node creation
  • Guard generation for tensor shapes and types
  • Variable tracking and symbolic execution

Let’s deep dive into each step separately!

🔬 Deep Dive: TorchDynamo’s Internal Architecture

Let’s dive deep into code base and check how it is tracing the program!

Entry Point

An execution of a Python program roughly consists of three stages:

  1. Initialization: handles import, etc.
  2. Compilation: run lexing, parsing, etc to convert python code into bytecode for execution.
  3. Interpretation: run CPython VM to execute the python function

The dynamo analysis was actually hacking into compilation stage, and rewrite the bytecode to be executed. When there is no already compiled function, we start a fresh analysis and tracing from scratch.

First step, it is using dis module to get the bytecode to be executed by CPython and run some initial cleanup of the Bytecode. You can check compiler here on details of how it is generating the bytecode. Also here is a great blog sharing how CPython is compiling and interpreting a python program

1
2
3
4
5
6

# torch version: 2.8.0+cu126

# From bytecode_transformation.py
def cleaned_instructions(code, safe=False) -> List[Instruction]:
    instructions = list(map(convert_instruction, dis.get_instructions(code)))
    # rest part are omitted

Symbolic Execution Engine

Instead of executing the bytecode via CPython interpreter, we created our own tracer to trace it. The core tracing happens in symbolic_convert.py through the InstructionTranslator class:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

# From symbolic_convert.py
class InstructionTranslatorBase:
    def __init__(...):
        #...
        output_graph = OutputGraph(...) # storing the fx graph, do code generation
        #...
        var = LazyVariableTracker.create(...) # create var tracker for input Tensors
    def step(self):
        """Process exactly one instruction, return False we should exit"""
        ip = self.instruction_pointer
        if ip is None:
            return False
        self.current_instruction = inst = self.instructions[ip]
        self.instruction_pointer = ip + 1
        
        # Dispatch to appropriate handler
        self.dispatch_table[inst.opcode](self, inst)

Instruction-by-Instruction Processing:

  • Each Python bytecode instruction has a corresponding handler
  • Variables are tracked as VariableTracker objects
  • Stack management maintains symbolic state

Bytecode Instruction Handlers

Each Python bytecode instruction has a specialized handler to generate output graph

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

# From symbolic_convert.py
def LOAD_FAST(self, inst):
    """Handle loading local variables"""
    self._load_fast(inst.argval)

def STORE_FAST(self, inst):
    """Handle storing local variables"""
    self._store_fast(inst.argval)

def CALL_FUNCTION(self, inst):
    """Handle function calls"""
    args = self.popn(inst.argval)
    fn = self.pop()
    self.call_function(fn, args, {})

def BINARY_ADD(self, inst):
    """Handle addition operations"""
    b, a = self.popn(2)
    self.push(a.call_method(self, "__add__", [b], {}))

VariableTracker

VariableTracker is the base class for tracked locals and stack values

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

class VariableTracker:
    def call_function(
        self,
        tx: "InstructionTranslator",
        args: Sequence["VariableTracker"],
        kwargs: "dict[str, VariableTracker]",
    ) -> "VariableTracker":
        # subclass is supposed to implement this ex. for "add" operation, it will be builtin functional variable tracker to insert and operation into graph located in Instruction Translator
        # ...
     def call_method(
        ...
    ) -> "VariableTracker":
        # handles method call for certain variable, ex. tensor.item; tensor[:, 1]
    def __init__(
        self,
        ...
    ) -> None:
        super().__init__()
        self.source = source # source for guard creation (ex. input of the graph)

OutputGraph

OutputGraph is class located in InstructionTranslator to manage the generated FX graph. When FX graph is traced, it will call compiles_and_call_fx_graph to further compile it into hardware-specific executables (ex. Torchinductor generating cuda kernels)

1
2
3
4
5
6
7
8
9

# simplified code
class OutputGraph:
    def __init__():
        self.export = export # whether to run in torch export (one graph) mode
        self.compiler_fn = compiler_fn # backend compiler ex. TorchInductor
        self.root_tx = root_tx # InstructionTranslator
    def compile_and_call_fx_graph():
        #Generate code from self.graph and return the Instruction()s to
        # call that generated code.

Guard Generation

When building up the graph, the VariableTracker generates comprehensive guards to make sure the pre-condition matches when we want to reuse an already compiled function

1
2
3
4

# From torch/_dyanmo/guards.py
def install_guard(self, guard_fn):
    # Install guards for tensor shapes, types, and values
    # Guards ensure graph validity across different inputs

Guards are validated before we call a compiled function. It is implemented in cpp (guard.cpp) to provide better performance

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

// example guard to validate tensor type match
class TYPE_MATCH : public LeafGuard {
 public:
  // type_id = id(type(obj))
  TYPE_MATCH(
      RootGuardManager* root_guard_manager,
      py::object type_id,
      py::object verbose_code_parts)
      : LeafGuard(root_guard_manager, std::move(verbose_code_parts)),
        _expected(py::cast<intptr_t>(std::move(type_id))) {}

  bool check_nopybind(PyObject* value) override { // borrowed ref
    // NOLINTNEXTLINE(performance-no-int-to-ptr)
    return Py_TYPE(value) == (void*)_expected;
  }

 private:
  // id of the type of the original object.
  intptr_t _expected;
};

See more details on how guard was implemented efficiently.

Graph Break Handling

When TorchDynamo encounters unsupported operations, it creates graph breaks:

1
2
3
4
5
6
7
8
9

# From symbolic_convert.py
def break_graph_if_unsupported(*, push):
    def decorator(inner_fn):
        def wrapper(self, inst):
            try:
                return inner_fn(self, inst)
            except Unsupported:
                # Create graph break and restart analysis
                self.current_speculation.fail_and_restart_analysis()

🎯 Torch Dynamo Summary

So far we have talked about overview of overall steps: byte-code execution, graph construction, graph-breaks, and guards. But the actual Torchdynamo is far more complicated than this overview (ex. important part is how AOTAutograd will work when we are tracing forward and backward pass, how do we handle mutations + functionalization, etc.) We will introduce some details on-demand in the following up series.

We now look into the phase 2 for hardware-specific optimization - TorchInductor

Phase 2: TorchInductor - Graph Optimization

What is TorchInductor?

TorchInductor is PyTorch’s backend compiler that takes the FX graph from TorchDynamo and optimizes it for execution on various backends (CPU, CUDA, etc.). For CUDA, it generates Triton kernels.

What is Triton?

Triton is a language and compiler for writing highly optimized GPU kernels. It provides:

  • High-level abstractions for GPU programming
  • Automatic optimization of memory access patterns
  • Parallel execution primitives
  • Integration with PyTorch

Interactive Example: See the Generation Process

After running the example in phase 1, check the debug directory:

1
2

# Find the latest debug directory
find torch_compile_debug/run_*/torchinductor/*/ -name "*.py" -o -name "*.txt" | head -10

This will show you the generated files including:

  • output_code.py: Generated Triton kernels
  • ir_pre_fusion.txt: IR before fusion
  • ir_post_fusion.txt: IR after fusion

🔬 Deep Dive: TorchInductor’s Triton Code Generation

Entry Point

First of all, when TorchDynamo has finished tracing the fx graph and reached the returning statement, it will trigger a compile_subgraph for further optimization

1
2
3
4
5
6
7
8
9
10
11
12
13

# in torch/_dynamo/symbolic_convert.py
# pytorch version: 2.8.0
class InstructionTranslator:
    #....
    def _return(self.inst):
        #...
        all_stack_locals_metadata = self.output.compile_subgraph(
            self,
            reason=GraphCompileReason(
                "return_value", [self.frame_summary()], graph_break=False
            ),
        )
        # ...

Then it goes into the entry point of inductor backend to compile the fx graph:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

# in torch/_inductor/compile_fx.py
# pytorch version: 2.8.0
def compile_fx(
    model_: GraphModule,
    example_inputs_: Sequence[InputType],
    inner_compile: Callable[..., OutputCode] = compile_fx_inner,
    config_patches: Optional[dict[str, Any]] = None,
    decompositions: Optional[dict[OpOverload, Callable[..., Any]]] = None,
    ignore_shape_env: bool = False,
) -> Union[Callable[[list[object]], Sequence[torch.Tensor]], str, list[str], Weights]:
    """
    Main entry point for compiling given FX graph.  Despite the fact that this
    lives in :mod:`torch._inductor`, this function is responsible for calling
    into AOT Autograd (and we will eventually get a callback to
    ``inner_compile`` to perform actual compilation.  In other words, this
    function orchestrates end-to-end compilation for the inductor backend when
    you use :func:`torch.compile`.

    NB: This function TAKES OWNERSHIP of the input ``model_`` and can potentially
    mutate it!  Make a copy if you need to preserve the original GraphModule.
    """

AOTAutograd

The AOTAutograd will basically help capturing backward graph if there are inputs that requires to capture gradient ahead of time (dispatcher code) (runtime wrapper to generate final function with compiled fw and bw).

Note: This module is not only embedded within Torchinductor, it is also used by torch.export, and can be the entry point of custom compiler backend.

See how does autograd work for torch.compile for high-level ideas.

1
2
3
4
5
6
7
8
9
10
11
12
13
14

# in torch/_inductor/compile_fx.py
# torch version: 2.8.0
def compile_fx(...):
        return aot_autograd(
                    fw_compiler=fw_compiler,
                    bw_compiler=bw_compiler,
                    inference_compiler=inference_compiler,
                    decompositions=decompositions,
                    partition_fn=partition_fn,
                    keep_inference_input_mutations=True,
                    cudagraphs=cudagraphs,
                    boxed_forward_device_index=forward_device,
                    ignore_shape_env=ignore_shape_env,
                )(model_, example_inputs_)

it will dispatch to the aot_dispatch_autograd function:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

# in torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py
# torch version: 2.8.0
def aot_dispatch_autograd(
    flat_fn,
    flat_args: list[Any],
    aot_config: AOTConfig,
    *,
    fw_metadata: ViewAndMutationMeta,
) -> DispatchReturn:
    """
    Autograd logic. Generates a joint graph, partitions it, manipulates the input with various wrappers,
    and returns a wrapped torch.autograd.Function with a forward and backward.
    """
    #...
    with dynamo_timed("aot_trace_joint_graph", log_pt2_compile_event=True):
        fx_g, joint_inputs, maybe_subclass_meta = aot_dispatch_autograd_graph(
            flat_fn, flat_args, aot_config, fw_metadata=fw_metadata
        ) # create joint graph
    #...
    fw_module, bw_module = aot_config.partition_fn(
                fx_g,
                joint_inputs,
                num_fwd_outputs=num_inner_fwd_outputs,
                static_lifetime_input_indices=static_lifetime_input_indices,
            )
    #...
    compiled_fn = post_compile(
        wrappers,
        compiled_fn,
        aot_config,
        runtime_metadata=fw_metadata,
    ) # run postprocessing to wrap the input and output of the graph

    return compiled_fn # return a torch.autograd.Function

when generating the joint graph, internally it is calling aot_dispatch_autograd_graph

1
2
3
4
5
6
7
8

def aot_dispatch_autograd_graph(..):
    #...
    joint_fn_to_trace, updated_joint_inputs = create_functionalized_fn(
        ...
    ) # functionalization to avoid mutation in the graph
    #...
    fx_g = _create_graph(joint_fn_to_trace, updated_joint_inputs, aot_config=aot_config) # this will create the joint graph, and will run "decomposition" to lower into a smaller set of supported operation

And internally it is running the fx graph and call fake mode of eager autograd (see details in here)

the “decomposition” is run in AOTAutograd as well, while all of the decomposition is registered in torch/_decom/decomposition.py (we should be able to register our own decomposition as well):

1
2
3
4
5
6
7
8
9

# in torch/_decomp/decomposition.py
# all of the decompositions are registered under this file

@register_decomposition(aten.silu)
@out_wrapper()
@pw_cast_for_opmath
def silu(self: Tensor) -> Tensor:
    return self * torch.sigmoid(self)

Also, The Joint Graph (graph module) will go through several passes for optimization. Ex. remove_noop, fuse_attention, etc. We can also register our own passes here.

Note: Here is a great debug point to showcase the when the graph is generated; Note: API: functorch.compile.draw_graph(graph_module) can be used to draw visualized diagram of the graph

The Joint fw+bw graph will be partitioned into forward and backward graph separately via max-flow / min-cut algorithm, which is targeting at minimize the tensor being saved for backward passes. See more details in this explanation. (reference of Ford–Fulkerson Algorithm)

Defined-by-run IR

After forward and backward graph is generated, it will generate a lowered defined-by-run IR for TorchInductor. These IR will have the layout information to execute the graph

1
2
3
4
5

# in torch/_inductor/compile_fx.py
class _InProcessFxCompile(FxCompile):
    def codegen_and_compile():
        #...
        graph = GraphLowering(...) # this will generate the defined-by-run IR

We can see the generated defined-by-run graph from ir_pre_fusion.txt

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

op0: SchedulerNode(ComputedBuffer)
op0.writes = [MemoryDep('buf0', c0, {c0: 1000})]
op0.unmet_dependencies = []
op0.met_dependencies = [MemoryDep('arg0_1', c0, {c0: 1000000}), MemoryDep('arg1_1', c0, {c0: 1000000})]
op0.outputs = [
    buf0: ComputedBuffer
    buf0.layout = FixedLayout('cuda:0', torch.float32, size=[1000], stride=[1])
    buf0.users = [NodeUser(node=SchedulerNode(name='op1'), can_inplace=False, is_weak=False)]
]
op0.group.device = cuda:0
op0.group.iteration = (1000, 1000)
op0.sizes = ([1000], [1000])
arg0_1_layout = FixedLayout('cuda:0', torch.float32, size=[1000, 1000], stride=[1000, 1])
arg1_1_layout = FixedLayout('cuda:0', torch.float32, size=[1000, 1000], stride=[1000, 1])
buf0_layout = FixedLayout('cuda:0', torch.float32, size=[1000], stride=[1])
class op0_loop_body:
    var_ranges = {p0: 1000, p1: 1000}
    index0 = 1000*p0 + p1
    index1 = p0
    def body(self, ops):
        get_index = self.get_index('index0')
        load = ops.load('arg0_1', get_index)
        get_index_1 = self.get_index('index0')
        load_1 = ops.load('arg1_1', get_index_1)
        add = ops.add(load, load_1)
        constant = ops.constant(2.0, torch.float32)
        mul = ops.mul(add, constant)
        relu = ops.relu(mul)
        reduction = ops.reduction(torch.float32, torch.float32, 'sum', relu)
        get_index_2 = self.get_index('index1')
        store_reduction = ops.store_reduction('buf0', get_index_2, reduction)
        return store_reduction


op1: SchedulerNode(ComputedBuffer)
op1.writes = [MemoryDep('buf1', 0, {})]
op1.unmet_dependencies = [MemoryDep('buf0', c0, {c0: 1000})]
op1.met_dependencies = []
op1.outputs = [
    buf1: ComputedBuffer
    buf1.layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
    buf1.users = [NodeUser(node=SchedulerNode(name='op2'), can_inplace=True, is_weak=False)]
]
op1.group.device = cuda:0
op1.group.iteration = (1, 1000)
op1.sizes = ([], [1000])
buf0_layout = FixedLayout('cuda:0', torch.float32, size=[1000], stride=[1])
buf1_layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
class op1_loop_body:
    var_ranges = {p0: 1000}
    index0 = p0
    index1 = 0
    def body(self, ops):
        get_index = self.get_index('index0')
        load = ops.load('buf0', get_index)
        reduction = ops.reduction(torch.float32, torch.float32, 'sum', load)
        get_index_1 = self.get_index('index1')
        store_reduction = ops.store_reduction('buf1', get_index_1, reduction)
        return store_reduction


op2: SchedulerNode(ComputedBuffer)
op2.writes = [MemoryDep('buf2', 0, {})]
op2.unmet_dependencies = [MemoryDep('buf1', 0, {})]
op2.met_dependencies = []
op2.outputs = [
    buf2: ComputedBuffer
    buf2.layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
    buf2.users = [NodeUser(node=OUTPUT, can_inplace=False, is_weak=False)]
]
op2.group.device = cuda:0
op2.group.iteration = (1, 1)
op2.sizes = ([], [])
buf1_layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
buf2_layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
class op2_loop_body:
    var_ranges = {}
    index0 = 0
    def body(self, ops):
        get_index = self.get_index('index0')
        load = ops.load('buf1', get_index)
        constant = ops.constant(1000.0, torch.float32)
        truediv = ops.truediv(load, constant)
        get_index_1 = self.get_index('index0')
        store = ops.store('buf2', get_index_1, truediv, None)
        return store

Scheduler

The next step is fusing the kernel via Scheduler, which is in charge of determining whether the op should be fused or not.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

# in torch/_inductor/scheduler.py

class Scheduler:
    def __init__(...):
        # ...
        log_ir_pre_fusion(self.nodes)
        self.num_orig_nodes = len(self.nodes)
        self.create_foreach_nodes()
        self.nodes = self.topological_sort_schedule(self.nodes)
        self.logged_slow_fusion = OrderedSet[tuple[str, str]]()
        if config._pre_fusion_custom_pass is not None:
            self.nodes = config._pre_fusion_custom_pass(self.nodes)
        self.nodes = self.fuse_nodes(self.nodes)
        if config._post_fusion_custom_pass is not None:
            self.nodes = config._post_fusion_custom_pass(self.nodes)
        self.merge_loops()
        self.finalize_multi_template_buffers()
        if config.combo_kernels:
            self.create_combo_kernel_nodes(num_ck_nodes=None)

        # Peak memory pass and overlap pass must run last, otherwise
        # other reordering passes could undo their effects.
        if config.reorder_for_peak_memory:
            from .memory import reorder_for_peak_memory

            self.nodes = reorder_for_peak_memory(
                self.nodes,
                self.name_to_buf,
                self.name_to_fused_node,
                OrderedSet(V.graph.graph_inputs.keys()),
                OrderedSet(V.graph.get_output_names()),
            )
        if config.reorder_for_compute_comm_overlap:
            self.nodes = comms.reorder_compute_and_comm_for_overlap(self.nodes)
        self.process_grouped_nodes()

        if torch._inductor.config.graph_partition:
            self.nodes = self.maybe_reorder_for_minimizing_partition(self.nodes)
            self.nodes = self.reorder_for_partition_with_simple_dependency(self.nodes)

        self.compute_last_usage()
        log_ir_post_fusion(self.nodes)
        # ...

We can see the post-fusion graph from ir_post_fusion.txt

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

#...
op1_op2: FusedSchedulerNode(SchedulerNode,SchedulerNode)
op1_op2.writes = [MemoryDep('buf1', 0, {}), MemoryDep('buf2', 0, {})]
op1_op2.unmet_dependencies = [MemoryDep('buf0', c0, {c0: 1000})]
op1_op2.met_dependencies = []
op1_op2.outputs = [
    buf1: ComputedBuffer
    buf1.layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
    buf1.users = [NodeUser(node=SchedulerNode(name='op2'), can_inplace=True, is_weak=False)]
    buf2: ComputedBuffer
    buf2.layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
    buf2.users = [NodeUser(node=OUTPUT, can_inplace=False, is_weak=False)]
]
op1_op2.snodes[0] =
op1: SchedulerNode(ComputedBuffer)
op1.writes = [MemoryDep('buf1', 0, {})]
op1.unmet_dependencies = [MemoryDep('buf0', c0, {c0: 1000})]
op1.met_dependencies = []
op1.outputs = [
    buf1: ComputedBuffer
    buf1.layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
    buf1.users = [NodeUser(node=SchedulerNode(name='op2'), can_inplace=True, is_weak=False)]
]
op1.group.device = cuda:0
op1.group.iteration = (1, 1000)
op1.sizes = ([], [1000])
buf0_layout = FixedLayout('cuda:0', torch.float32, size=[1000], stride=[1])
buf1_layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
class op1_loop_body:
    var_ranges = {p0: 1000}
    index0 = p0
    index1 = 0
    def body(self, ops):
        get_index = self.get_index('index0')
        load = ops.load('buf0', get_index)
        reduction = ops.reduction(torch.float32, torch.float32, 'sum', load)
        get_index_1 = self.get_index('index1')
        store_reduction = ops.store_reduction('buf1', get_index_1, reduction)
        return store_reduction
op1_op2.snodes[1] =
op2: SchedulerNode(ComputedBuffer)
op2.writes = [MemoryDep('buf2', 0, {})]
op2.unmet_dependencies = [MemoryDep('buf1', 0, {})]
op2.met_dependencies = []
op2.outputs = [
    buf2: ComputedBuffer
    buf2.layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
    buf2.users = [NodeUser(node=OUTPUT, can_inplace=False, is_weak=False)]
]
op2.group.device = cuda:0
op2.group.iteration = (1, 1)
op2.sizes = ([], [])
buf1_layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
buf2_layout = FixedLayout('cuda:0', torch.float32, size=[], stride=[])
class op2_loop_body:
    var_ranges = {}
    index0 = 0
    def body(self, ops):
        get_index = self.get_index('index0')
        load = ops.load('buf1', get_index)
        constant = ops.constant(1000.0, torch.float32)
        truediv = ops.truediv(load, constant)
        get_index_1 = self.get_index('index0')
        store = ops.store('buf2', get_index_1, truediv, None)
        return store

Triton Codegen

The last step will be the triton codegen

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

# in torch/_inductor/codegen/simd.py
class SIMDScheduling(BaseScheduling):
    def codegen_node(
            self, node: Union[scheduler.FusedSchedulerNode, scheduler.SchedulerNode]
        ):`
        """
        Given a set of pre-fused nodes, generate a Triton kernel.
        """

        nodes: list[scheduler.SchedulerNode] = node.get_nodes()  # type: ignore[assignment]

        if torch._inductor.config.triton.coalesce_tiling_analysis:
            coalesce_analysis = analyze_memory_coalescing(node)
        else:
            coalesce_analysis = None
        _, (numel, rnumel) = max(nodes, key=lambda x: int(x.is_reduction())).group

        node_schedule = self.generate_node_schedule(nodes, numel, rnumel)
        schedule_log.debug("Schedule:\n %s", node_schedule)

        return self.codegen_node_schedule(
            SIMDKernelFeatures(node_schedule, numel, rnumel, coalesce_analysis)
        )

We will be able to see the final generated triton code in output_code.py

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

#...
@triton_heuristics.persistent_reduction(
    size_hints={'x': 1024, 'r0_': 1024},
    reduction_hint=ReductionHint.INNER,
    filename=__file__,
    triton_meta={'signature': {'in_ptr0': '*fp32', 'in_ptr1': '*fp32', 'out_ptr0': '*fp32', 'xnumel': 'i32', 'r0_numel': 'i32'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=132, cc=90, major=9, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_per_fused_add_mul_relu_silu_sum_0', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': True, 'num_load': 2, 'num_reduction': 1, 'backend_hash': 'F0EE7D6F735E9CCA02E1254E65809C6B6C17DFB7C86025DA2A9F7ED74AE9DFB2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'tiling_scores': {'x': 8000, 'r0_': 8000000}}
)
@triton.jit
def triton_per_fused_add_mul_relu_silu_sum_0(in_ptr0, in_ptr1, out_ptr0, xnumel, r0_numel):
    xnumel = 1000
    XBLOCK: tl.constexpr = 1
    r0_numel = 1000
    R0_BLOCK: tl.constexpr = 1024
    rnumel = r0_numel
    RBLOCK: tl.constexpr = R0_BLOCK
    xoffset = tl.program_id(0) * XBLOCK
    xindex = tl.full([1], xoffset, tl.int32)
    xmask = tl.full([R0_BLOCK], True, tl.int1)
    r0_index = tl.arange(0, R0_BLOCK)[:]
    r0_offset = 0
    r0_mask = r0_index < r0_numel
    roffset = r0_offset
    rindex = r0_index
    r0_1 = r0_index
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (r0_1 + 1000*x0), r0_mask, other=0.0)
    tmp1 = tl.load(in_ptr1 + (r0_1 + 1000*x0), r0_mask, other=0.0)
    tmp2 = tmp0 + tmp1
    tmp3 = 2.0
    tmp4 = tmp2 * tmp3
    tmp5 = tmp4 + tmp0
    tmp6 = tl.full([1], 0, tl.int32)
    tmp7 = triton_helpers.maximum(tmp6, tmp5)
    tmp8 = tl.sigmoid(tmp7)
    tmp9 = tmp7 * tmp8
    tmp10 = tl.broadcast_to(tmp9, [R0_BLOCK])
    tmp12 = tl.where(r0_mask, tmp10, 0)
    tmp13 = triton_helpers.promote_to_tensor(tl.sum(tmp12, 0))
    tl.store(out_ptr0 + (x0), tmp13, None)
''', device_str='cuda')

🎓 Key Takeaways

What You’ve Learned

  1. TorchDynamo: Converts Python functions to computational graphs through sophisticated bytecode analysis and symbolic execution
  2. TorchInductor: Optimizes graphs and generates backend-specific code (Triton)

Next Steps

In the next series, we will:

Deep Dive into Complicated Cases:

We will deep dive further into some common case handled by torch compile when we are running model training, including but not limited to:

  • AOT Autograd
  • torch.compile and backward pass
  • torch.comiple and Distributed Data Parallel (DDP)
  • torch.compile and Fully Sharded Distributed Parallel (FSDP2)

Practical Debugging Guide

We will share more practical debugging tools and variables when dealing with common issues.

From Triton kernel to GPU binary

What happens next after Triton Python DSL? We will have further compliation to optimize specifically for certain hardware (ex. tritonIR -> LLVM IR -> PTX -> cubin). We will introduce these in the following sections as well!

Machine Learning Performance, torch.compile
pytorch torch.compile triton optimization deep-learning
This post is licensed under CC BY 4.0 by the author.