Computational Graphs as an Intermediate Representation

A serialized inference model is a flat table. One array lists nodes, each naming an operator, the value identifiers it reads, the value identifiers it writes, and a small attribute dictionary holding kernel size, padding, epsilon, or axis. A second array describes every value with a dtype and a shape. Nothing else survives the export: no instruction pointer, no branch, no variable that can be assigned twice, no call stack. That austerity is deliberate. Every transformation applied to a compiled model, from operator fusion to memory planning, is a rewrite of that table, and each rewrite is legal only because the table has properties an ordinary program does not.

The dataflow formulation is old. Dennis described programs as directed graphs of operators consuming tokens from arcs in 1974, precisely because the absence of a program counter makes the available parallelism explicit. Machine learning compilers rediscovered the same representation for a narrower reason: a trained model is a fixed composition of tensor functions with no loops over data-dependent conditions, which is exactly the subset of computation a dataflow graph can express.

Each value is produced exactly once

The single most consequential invariant is that every value has at most one defining node. This is the static single assignment property, and it is what makes the graph tractable. Liveness needs no dataflow fixpoint: a value is born at its defining node and dies at its last consumer in the chosen execution order. Dead code elimination is a reachability query from the graph outputs backwards. Rewriting a node’s inputs cannot invalidate any other node’s definition, so a pass can transform one operator at a time without a global rename step. In a program with mutable variables, the same edits would require recomputing reaching definitions after every change.

The invariant also makes cost models composable. A pass that fuses two nodes knows the fused node produces the same values, because there is no other writer to worry about and no intervening assignment that could be observed.

There is one crack in the abstraction: aliasing. reshape, transpose, slice, and broadcast can be represented either as copies or as views onto the same storage. If the graph records a view as an ordinary value, an optimizer may believe it has two independent buffers and reorder a write past a read that actually observes it. Real graph formats therefore carry either an explicit alias annotation on view-like operators, or a rule that such operators never alias and the runtime materialises a copy. StableHLO takes the second route for most operations by defining them functionally, which costs a copy at the boundary of a view chain but keeps every optimization local. Frameworks that expose in-place mutation of a tensor to the graph must record the mutation as an ordering dependency, which reintroduces exactly the analysis the SSA form was meant to avoid.

Topological order is the entire schedule

Any linear extension of the partial order defined by the edges is a correct execution order. Producing one is Kahn’s algorithm: count the in-degree of each node, start from the nodes with no unproduced inputs, and emit a node once all of its producers have been emitted. The cost is $O(V + E)$ with $V$ nodes and $E$ edges, and a failure to emit all $V$ nodes proves the graph has a cycle.

from dataclasses import dataclass, field

@dataclass
class Node:
    op: str
    inputs: list[int]
    outputs: list[int]
    attrs: dict = field(default_factory=dict)

def topological_order(nodes: list[Node]) -> list[int]:
    producer = {v: i for i, n in enumerate(nodes) for v in n.outputs}
    indeg = [0] * len(nodes)
    consumers: list[list[int]] = [[] for _ in nodes]
    for i, n in enumerate(nodes):
        for v in n.inputs:
            if v in producer:            # graph inputs have no producer
                indeg[i] += 1
                consumers[producer[v]].append(i)
    ready = [i for i, d in enumerate(indeg) if d == 0]
    order: list[int] = []
    while ready:
        i = ready.pop()                  # LIFO: prefer the most recently freed node
        order.append(i)
        for c in consumers[i]:
            indeg[c] -= 1
            if indeg[c] == 0:
                ready.append(c)
    if len(order) != len(nodes):
        raise ValueError("cycle in graph")
    return order

The choice among the valid orders is not cosmetic. It determines how long each intermediate stays live, and therefore the peak size of the memory arena. Consider a residual block whose skip branch is computed early and added at the end. Emitting the skip branch at its definition keeps its output live across every node in between. Emitting it immediately before the add, which is equally legal, shortens that interval to one step. The same freedom lets a scheduler interleave independent branches to keep a device busy, at the cost of holding both branches’ intermediates simultaneously. Choosing a linear extension is what scheduling means here, and the two objectives, peak memory and utilisation, pull in opposite directions.

Cycles appear in two places. Recurrent models unrolled over time produce acyclic graphs by construction, with the state carried as an explicit loop variable. Graphs that retain a loop, such as ONNX Loop or StableHLO while, hide the cycle inside a region node whose body is itself a graph. An autoregressive decoder usually keeps the token loop outside the compiled graph entirely, launching the same graph once per step, which is why capturing the launch sequence becomes relevant for that workload.

Shape inference must run before anything else

Every pass downstream assumes it knows the byte size of each value, so shape inference is the first thing that runs after import. Each operator carries a shape function: elementwise operations preserve the input shape after broadcasting, a matrix multiply contracts the inner dimension, a reduction removes or keeps an axis of size one depending on a flag, a convolution computes output extent from input extent, kernel size, padding, and stride. Broadcasting follows the right-aligned rule: two shapes are compatible if, for each trailing axis, the extents are equal or one of them is one.

Shape inference also fixes dtype promotion. Mixing a float32 activation with a float16 weight either promotes to float32 or requires an explicit cast inserted by the importer; leaving that decision to the runtime would make kernel selection depend on values rather than on the graph.

The failure mode is silent rather than loud. A shape function that is wrong by one axis does not crash the compiler. It produces a memory plan with a buffer that is too small, a kernel that vectorises over the wrong dimension, or a fused group whose loop bounds do not agree. The corruption surfaces as a wrong result much later. This is why shape inference is usually validated by a reference interpreter that executes the graph one operator at a time on small inputs, and why formats that publish a specification of each operator’s shape semantics, as StableHLO does, are easier to compile correctly than formats that leave the semantics to the exporter.

What the representation cannot express

A graph says nothing about control flow that depends on data. An operator whose output shape is a function of its input values, such as nonzero, unique, or a top-k with a runtime k, cannot be given a static shape at compile time. The runtime must read a value back to the host, synchronise, allocate, and then continue. Each such readback is a pipeline stall, and it also breaks the assumption that the whole graph can be captured and replayed as a fixed sequence of launches.

The representation also cannot express host-side effects, random number generation without an explicit seed value threaded through the graph, or communication with an external service. A model that calls into Python between layers is not one graph; it is a sequence of graphs separated by breaks, each of which is compiled and optimised in isolation.

Finally, the graph carries no memory layout beyond dtype and logical shape. Whether a four-dimensional activation is stored channels-first or channels-last is a decision made during layout assignment, and it is invisible in the exported model. Two runtimes given the same file may therefore reach very different performance without either being wrong.

Where the abstraction costs memory

A bytecode interpreter for the same computation holds intermediates on a stack and overwrites them as soon as a value is consumed, so its working set is bounded by expression depth. A naive graph executor allocates a distinct buffer per value and frees each at its last use, which is a larger working set whenever the topological order is not chosen to minimise liveness. The graph representation therefore starts at a disadvantage in memory and has to earn it back: through fusion, which removes intermediate buffers entirely, and through planning, which overlaps the remaining ones in a single arena.

The compensating advantage is that the compiler can see the whole computation before anything executes. A stack machine cannot reorder its instructions; a graph can be reordered, fused, and specialised per shape. Every later transformation is a way of spending that advantage.