ATen, core ATen, prims

make_fx traces a function into an FX graph and takes a decomposition table as an argument, which makes it the cleanest instrument for this. Trace torch.nn.functional.gelu three times, changing only the table, and you get three graphs of the same function at three depths.

Plain, it is one node: aten.gelu.default. With the core table, it is still one node, because gelu is already in the core set and the table has no rule that touches it. Under TorchRefsMode, which routes every torch call through the reference implementations in torch._refs, it becomes five prims nodes and the arithmetic is finally visible: multiply by a half, multiply by one over the square root of two, take the error function, add one, multiply the two halves together.

The constant 0.7071067811865476 in that printout is the whole point of the bottom rung. At the prims level nothing is named after a neural network idea any more. There are only elementwise operations on shaped buffers, which is a small enough vocabulary that a new hardware backend can implement all of it.

import torch
from torch.fx.experimental.proxy_tensor import make_fx
from torch._decomp import core_aten_decompositions
from torch._prims.context import TorchRefsMode

def f(x):
    return torch.nn.functional.gelu(x)

x = torch.randn(3)
print("aten:")
print(make_fx(f)(x).code.strip())
print("core aten:")
print(make_fx(f, decomposition_table=core_aten_decompositions())(x).code.strip())
print("prims:")
with TorchRefsMode():
    print(make_fx(f)(x).code.strip())
print("core_aten_decompositions entries:", len(core_aten_decompositions()))

The output of that script, verbatim, is below. Read the last line first: the core decomposition table on this build holds 373 entries, which is the number of rewrite rules, not the size of the core set itself.

Set that against the numbers from the last lesson. There are 817 ATen packets carrying 2124 overloads, and 132 prims packets. A backend author facing 2124 signatures has an impossible job; facing 132 has a finite one, at the cost of writing every fused kernel out of pieces. Core ATen sits between the two, small enough to implement and still coarse enough to keep the fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → opportunities the prims level has already thrown away.

Which rung you want depends on what you are building. A compiler that fuses aggressively wants ops small enough to rearrange. A vendor library with a hand-written attention kernel wants the op left whole so it can pattern-match it. Both of those preferences are expressed by choosing a decomposition table, not by changing the model.

aten:
def forward(self, x_1):
    gelu = torch.ops.aten.gelu.default(x_1);  x_1 = None
    return gelu
core aten:
def forward(self, x_1):
    gelu = torch.ops.aten.gelu.default(x_1);  x_1 = None
    return gelu
prims:
def forward(self, x_1):
    mul = torch.ops.prims.mul.default(x_1, 0.5)
    mul_1 = torch.ops.prims.mul.default(x_1, 0.7071067811865476);  x_1 = None
    erf = torch.ops.prims.erf.default(mul_1);  mul_1 = None
    add = torch.ops.prims.add.default(erf, 1.0);  erf = None
    mul_2 = torch.ops.prims.mul.default(mul, add);  mul = add = None
    return mul_2
core_aten_decompositions entries: 373

Four ops through the same comparison show the shape of the boundary better than any list. silu is not core, so the table rewrites it into sigmoid plus a mul. hardswish is not core either, and comes out as add, two clamps, a mul and a div. _softmax and native_layer_norm both survive untouched, because both are in the core set and a backend is expected to implement them.

The fourth case is the interesting one, and it is not about neural network ops at all. nn.functional.linear traces to aten.t.default followed by addmm. Under the core table, the addmm is unchanged and the t becomes aten.permute.default with an explicit [1, 0]. Same node count, same work, one fewer signature to implement.

That is the rule the core set is built on. Decomposition is not simplification and it is not always a rewrite into more nodes. It is a reduction in the number of distinct signatures a backend has to understand, and sometimes the cheapest way to reach it is to replace a convenience op with the general one it is a special case of.

You do not have to build the trace by hand to get this. An ExportedProgram carries run_decompositions(), which returns a new program with the table applied, and running it on a linear layer followed by a silu shows both rewrites from the table above happening at once on real graph.

Three things change in that printout and only two of them are decompositions. t becomes permute with an explicit permutation. silu becomes sigmoid and mul, and the addmm node picks up a second user because both of them consume it. The third change is cosmetic: the user input, printed as %l_x_ before, comes back as %arg2_1, because the pass renumbers every placeholder into one sequence.

Chapter 10 says an exported program reaches the bridges decomposed toward a core set of ATen ops. This is the call that does it, and the reason the bridge cares is the same reason a backend author cares: on the other side of the crossing there is a compiler whose op list is finite and does not include silu.

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.nn.functional.silu(self.fc(x))

torch.manual_seed(0)
ep = torch.export.export(Net(), (torch.randn(2, 3),))
print(ep.graph)
print(ep.run_decompositions().graph)

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.nn.functional.silu(self.fc(x))

torch.manual_seed(0)
ep = torch.export.export(Net(), (torch.randn(2, 3),))
print(ep.graph)
print(ep.run_decompositions().graph)

# graph():
#     %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
#     %arg1_1 : [num_users=1] = placeholder[target=arg1_1]
#     %l_x_ : [num_users=1] = placeholder[target=l_x_]
#     %t : [num_users=1] = call_function[target=torch.ops.aten.t.default](args = (%arg0_1,), kwargs = {})
#     %addmm : [num_users=1] = call_function[target=torch.ops.aten.addmm.default](args = (%arg1_1, %l_x_, %t), kwargs = {})
#     %silu : [num_users=1] = call_function[target=torch.ops.aten.silu.default](args = (%addmm,), kwargs = {})
#     return (silu,)
# graph():
#     %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
#     %arg1_1 : [num_users=1] = placeholder[target=arg1_1]
#     %arg2_1 : [num_users=1] = placeholder[target=arg2_1]
#     %permute : [num_users=1] = call_function[target=torch.ops.aten.permute.default](args = (%arg0_1, [1, 0]), kwargs = {})
#     %addmm : [num_users=2] = call_function[target=torch.ops.aten.addmm.default](args = (%arg1_1, %arg2_1, %permute), kwargs = {})
#     %sigmoid : [num_users=1] = call_function[target=torch.ops.aten.sigmoid.default](args = (%addmm,), kwargs = {})
#     %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%addmm, %sigmoid), kwargs = {})
#     return (mul,)

The jax path solves the same problem with the same move and different names, which is worth noticing because it tells you the problem is structural rather than a PyTorch quirk. There, a client-level op either has an HLO counterpart and survives, or it falls through a table of decomposition patterns into spec ops. The lesson on how a jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → becomes HLO works one such table line by line.

The difference is where the table lives. PyTorch's decompositions are written in Python, in torch._decomp and torch._refs, and you can swap the table per compilation. The XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → equivalent is compiled into the compiler. That is why a torch backend can be selective about which ops it wants left whole, and why the same model can reach two backends with two different op sets from one export.

A practical consequence for reading profiles: when a fused kernel you expected does not appear, check whether the op survived decomposition before you go looking at the codegen. An op that got rewritten into five pieces upstream never reached the fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → pass as one thing.