Torch Internals (Part 4) - TorchInductor

In the last 3 blog posts, we have understood fx graphs, how TorchDynamo captures an FX graph and how AOTAutograd splits it into forward/backward graphs. The final piece of the compile stack is TorchInductor which is the compiler that takes those FX graphs and turns them into efficient Triton kernels.

Inductor works in three phases. First, it lowers the FX graph into its own IR (a recipe for computing each element). Then it schedules and fuses those IR nodes. Finally, it walks through the IR and generates Triton code.

Let's understand the IR

The first thing inductor does is to lower the FX graph produced by TorchDynamo into its own IR. So it's very important to try and understand what it looks like.

@torch.compile
def f(x):
    b = torch.floor(x) + torch.ceil(x)
    c = b.sum(dim=-1)
    d = c + 1
    return d

For the above function, let's see how each line gets lowered into an IR.

1. Placeholders/Inputs

The input x is converted into an IR that looks like this -

TensorBox(
    StorageBox(
        InputBuffer(
            name="arg0",
            layout=FixedLayout(
                device="cuda:0",
                dtype=torch.float32,
                size=[32, 512, 1024],
                stride=[524288, 1024, 1]
            )
        )
    )
)

InputBuffer stores metadata about the input tensor like device, dtype, shape etc but no computation happens here. Think of it as a reference to data that already exists in memory.

2. Pointwise Operations

torch.floor(x) becomes:

TensorBox(
    StorageBox(
        Pointwise(
            device='cuda:0',
            dtype=torch.float32,

            def inner_fn(index):
                i0, i1, i2 = index
                tmp0 = ops.load(arg0_1, i2 + 1024*i1 + 524288*i0)
                tmp1 = ops.floor(tmp0)
                return tmp1

            ranges=[32, 512, 1024]
        )
    )
)

For pointwise operations, inductor stores a recipe for computing one output element but it never actually computes it. The function that holds this recipe is called the inner_fn.

You can read the function and see the logic to load the element across the stride and compute floor.

3. Pointwise Fusion

Cool, that was for pointwise operations. What about pointwise fusion? For example floor(x) + ceil(x)

TensorBox(
    StorageBox(
        Pointwise(
            def inner_fn(index):
                i0, i1, i2 = index

                tmp0 = ops.load(arg0_1, i2 + 1024*i1 + 524288*i0)
                tmp1 = ops.floor(tmp0)

                tmp2 = ops.load(arg0_1, i2 + 1024*i1 + 524288*i0)
                tmp3 = ops.ceil(tmp2)

                tmp4 = tmp1 + tmp3
                return tmp4
        )
    )
)

Ah, instead of having individual inner_fn, both of them got rolled into a single function.

4. Reduction Operations

Reduction operations are expressed in a different way since unlike pointwise operators, they depends on many other elements.

TensorBox(
    StorageBox(
        ComputedBuffer(
            name="buf0",
            data=Reduction(
                ranges=[32, 512],
                reduction_ranges=[1024],

                def inner_fn(index, rindex):
                    i0, i1 = index
                    r0 = rindex

                    tmp0 = ops.load(arg0_1, r0 + 1024*i1 + 524288*i0)
                    tmp1 = ops.floor(tmp0)

                    tmp2 = ops.load(arg0_1, r0 + 1024*i1 + 524288*i0)
                    tmp3 = ops.ceil(tmp2)

                    tmp4 = tmp1 + tmp3
                    return tmp4
            )
        )
    )
)

Conceptually this is:

for i0 in range(32):
    for i1 in range(512):
        total = 0
        for r0 in range(1024):
            total += floor(x[i0,i1,r0]) + ceil(x[i0,i1,r0])
        output[i0,i1] = total

Remember we are trying to express this line c = b.sum(dim=-1), so we have the floor, ceil, and add computations appear again inside the reduction. This is not recomputation rather those recipes are simply copied into this reduction recipe.

Also notice, the reduction result is wrapped in a ComputedBuffer(name="buf0"). This means the result is now a logical tensor that future operations can read. It does not mean memory has been allocated and whether it lives in registers, shared memory, or global memory is decided later by the scheduler.

The last pointwise operation

The last line in the function c + 1 becomes -

TensorBox(
    StorageBox(
        Pointwise(
            def inner_fn(index):
                i0, i1 = index
                tmp0 = ops.load(buf0, i1 + 512*i0)
                tmp1 = ops.constant(1, torch.float32)
                tmp2 = tmp0 + tmp1
                return tmp2
        )
    )
)

Unlike the earlier pointwise nodes, this one loads from buf0 (the reduction result) instead of the input tensor.

The output node wraps the final computation inside another ComputedBuffer(name="buf1") and it is the final tensor returned by the compiled function.

How Lowering Happens

Great! We have now seen some of the common IRs for some important operations. Let's now look at how lowering actually happens.

We've seen that inner_fn is a Python function full of ops.load, ops.floor, etc. But these aren't really pytorch ops, rather they're abstract ops that gets replaced depending on what handler is installed during codegen.

ops.* like ops.floor, ops.load, ops.add are defined on a class called OpOverrides. Take a look at this -

class TritonOverrides(OpOverrides):
    @staticmethod
    def floor(x):
        return f"tl.math.floor({x})"

    @staticmethod
    def ceil(x):
        return f"tl.math.ceil({x})"

    @staticmethod
    def add(x, y):
        return f"({x} + {y})"

    @staticmethod
    def load(name, index):
        return f"tl.load({name} + ({index}), None)"

So basically what happens is when the inner_fn is executed during codegen, they return these hardcoded code strings but they don't compute anything.

Here's an example.

def inner_fn(index):
    tmp0 = ops.load("in_ptr0", "x0")       # "tl.load(in_ptr0 + (x0), None)"
    tmp1 = ops.floor(tmp0)                 # "tl.math.floor(tl.load(...))"
    tmp2 = ops.ceil(tmp0)                  # "tl.math.ceil(tl.load(...))"
    tmp3 = ops.add(tmp1, tmp2)             # "(tl.math.floor(...) + tl.math.ceil(...))"
    return tmp3

Cool, but where did these intermediate tmp0, tmp1 etc come from?

Common Subexpression Elimination

If you remember in the IR, inner_fn loads the same input twice (once for floor, once for ceil). But in the generated Triton code, it loads this only once.

tmp0 = tl.load(in_ptr0 + (x0), None)
tmp1 = tl.math.floor(tmp0)
tmp2 = tl.math.ceil(tmp0)          # reuses tmp0, no second load

This happens because Inductor's CSE class deduplicates identical expressions:

class CSE:
    def generate(self, buffer, expr):
        var = self.cache.get(expr)
        if not var:
            var = self.newvar()
            self.cache[expr] = var
            buffer.writeline(f"{var} = {expr}")
        return var

The two loads have the same address expression, so the second one returns the same tmp0 variable.