Reading a real graph

FX has a fixed vocabulary of node opcodes, and there are only six. A module with one linear layer, one buffer, one operator, and two method calls produces five of them, and the sixth is the return. That is the whole grammar, which is why the same reading habit works on a dynamo capture, an aot forward, and an exported program.

torch.fx.symbolic_trace is the oldest and simplest way to get one. It runs your forward with proxy objects and records what happens, which is not what dynamo does and not what export does, but it produces the same node types, and it leaves modules unopened so you can see call_module and get_attr in the same picture.

Read the printout below top to bottom. %x is the input. %fc calls a submodule by name. %shift fetches a registered buffer off the module, no computation at all. %add calls operator.add, the Python built-in behind the +. %relu and %sum_1 are method calls on the tensor that came before. The last line is the return.

import torch
import torch.nn as nn
from collections import Counter

class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(3, 4)
        self.register_buffer("shift", torch.zeros(4))
    def forward(self, x):
        return (self.fc(x) + self.shift).relu().sum()

gm = torch.fx.symbolic_trace(M())
print(gm.graph)
print(Counter(n.op for n in gm.graph.nodes))

import torch
import torch.nn as nn
from collections import Counter

class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(3, 4)
        self.register_buffer("shift", torch.zeros(4))
    def forward(self, x):
        return (self.fc(x) + self.shift).relu().sum()

gm = torch.fx.symbolic_trace(M())
print(gm.graph)
print(Counter(n.op for n in gm.graph.nodes))

# graph():
#     %x : [num_users=1] = placeholder[target=x]
#     %fc : [num_users=1] = call_module[target=fc](args = (%x,), kwargs = {})
#     %shift : [num_users=1] = get_attr[target=shift]
#     %add : [num_users=1] = call_function[target=operator.add](args = (%fc, %shift), kwargs = {})
#     %relu : [num_users=1] = call_method[target=relu](args = (%add,), kwargs = {})
#     %sum_1 : [num_users=1] = call_method[target=sum](args = (%relu,), kwargs = {})
#     return sum_1
# Counter({'call_method': 2, 'placeholder': 1, 'call_module': 1, 'get_attr': 1, 'call_function': 1, 'output': 1})

Every non-return line has the same five parts, and the printout spaces them out the same way each time: %name : [num_users=N] = opcode[target=T](args = (...), kwargs = {...}). The name is the SSA-style handle other lines refer to. The opcode is one of the six. The target is what the opcode dispatches on.

num_users is the one people skim past, and it is the field that tells you whether a value is about to be freed. It counts how many later nodes take this node as an argument, which is why the gm.code view of the same graph is full of ; x = None statements: the codegen drops the reference as soon as the last user has consumed it, so the buffer can be released mid-graph.

It is also the field a graph pass reads to decide whether an edit is legal. Dead code elimination is exactly the rule that a node with no users and no side effects can go, and you can see the result rather than the rule in the min-cut forward from lesson one, where the values the backward stopped needing simply do not appear in the return list.

The .default that ends almost every ATen target in these graphs is not punctuation. torch.ops.aten.sum is an overload packet, a family of C++ signatures sharing one name, and the last component picks one member of that family. Ask the packet directly and it lists them.

So aten.sum.default is the reduce-everything signature and aten.sum.dim_IntList is the one taking a dimension list, and the aot backward in lesson one used the second because the bias gradient sums over dimension 0 only. A graph that names a specific overload has already resolved which C++ kernel it means, which is why an FX graph at this level is unambiguous in a way your Python source is not.

Counting the surface makes the scale concrete. On this build, torch.ops.aten exposes 817 packets, which between them carry 2124 overloads. Seven of the packets do not report overloads at all. torch.ops.prims exposes 132. Those three numbers are the reason decompositions exist, which is the next lesson.

import torch

packets = dir(torch.ops.aten)
total = 0
skipped = 0
for name in packets:
    try:
        total += len(getattr(torch.ops.aten, name).overloads())
    except Exception:
        skipped += 1
print("aten packets:", len(packets), "overloads:", total, "packets with none:", skipped)
print("prims packets:", len(dir(torch.ops.prims)))
print("aten.sum overloads:", torch.ops.aten.sum.overloads())

# aten packets: 817 overloads: 2124 packets with none: 7
# prims packets: 132
# aten.sum overloads: ['dim_IntList', 'default', 'dim_DimnameList', 'DimnameList_out', 'IntList_out', 'out', 'int', 'float', 'complex', 'bool']

The printed graph shows no shapes and no dtypes, which makes it look like less information than it holds. Every node carries a meta dict, and for graphs that went through fake-tensor tracing the key val holds a tensor with the right shape, dtype and device and no storage behind it. Reading that dict turns a graph dump into a shape table.

This is the same shape-propagation you would otherwise do by hand while staring at a matmul, done once during tracing and kept. Note where the t node lands in the table below: (4, 3) in, (3, 4) out, which is the transpose an nn.Linear needs before addmm can consume the weight.

One caution on the placeholder names in that output. arg0_1 and l_x_ are what torch 2.2.2 prints; current torch names lifted parameters after their module path instead, so an exported graph there begins with something like p_fc_weight. The structure is the same and the names are not, so re-run this on your own build before you quote a name.

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(3, 4)
    def forward(self, x):
        return torch.relu(self.fc(x)).sum()

torch.manual_seed(0)
ep = torch.export.export(Net(), (torch.randn(2, 3),))
for n in ep.graph.nodes:
    v = n.meta.get("val")
    shape = tuple(v.shape) if isinstance(v, torch.Tensor) else v
    print(f"{n.op:13s} {n.name:9s} {str(n.target)[:22]:22s} {shape}")

# placeholder   arg0_1    arg0_1                 (4, 3)
# placeholder   arg1_1    arg1_1                 (4,)
# placeholder   l_x_      l_x_                   (2, 3)
# call_function t         aten.t.default         (3, 4)
# call_function addmm     aten.addmm.default     (2, 4)
# call_function relu      aten.relu.default      (2, 4)
# call_function sum_1     aten.sum.default       ()
# output        output    output                 None

Put the printouts from this lesson and the last one side by side and a pattern falls out that is worth carrying as a reading rule. The deeper the graph, the fewer opcodes it uses. The symbolic trace above holds all six. The dynamo capture in lesson one used call_module and call_method with no get_attr at all, because its parameters were still inside an unopened Linear. The aot and exported graphs have none of those three: every computing node is a call_function on an OpOverload.

That narrowing is what makes the bottom of the stack compilable. A backend that has to understand call_module has to understand arbitrary Python classes. A backend that only sees call_function on resolved ATen overloads has a finite list to implement, and the size of that list is the whole subject of the next lesson.

The narrowing also tells you where you are when someone hands you a graph with no context. Count the opcodes present. If you see a module call, you are above aot_autograd; if every line is an overload, you are below it.