Constant Folding and Algebraic Rewrites on a Fixed Graph

A trained model arrives carrying weight tensors, and the graph around them contains a surprising amount of arithmetic that does not depend on the input. A batch-norm layer in inference mode is an affine map whose scale and offset can be computed once from the stored statistics. A reshape of a reshaped tensor is one reshape. A transpose of a transpose is nothing. A concatenation of two constant slices of the same weight tensor is a constant. Each of these is a rewrite that removes nodes, and each rewrite is only worth doing because it shrinks everything downstream: fewer launches, fewer intermediates, and a shorter list of candidates for fusion.

The mechanism for the first category is simple and the mechanism for the second is not, and the difference between them is where most of the difficulty in this pass lives.

Evaluating the constant subgraph

Constant folding marks every node whose inputs are all constants, executes it with the same kernel the runtime would use, and replaces its outputs with constant values. Doing it in one sweep is wrong, because folding a node can make its consumers constant. The pass iterates to a fixpoint: fold, then re-examine consumers, until no node changes. A worklist implementation pushes each consumer of a newly folded value, which bounds the work to the number of edges rather than the number of nodes times the graph depth.

The subtlety is that folding must produce bit-identical results to the unfolded graph, or the compiler has changed the model. Two conditions threaten that. First, the kernel used at compile time may differ from the one selected at run time: a compiler running on the host CPU and folding an operation in double precision, then emitting the result as a float32 constant, produces a different value than a device kernel computing in float32. The fix is to fold in the dtype of the graph, using a reference implementation of the operator’s semantics rather than the fastest available routine.

Second, folding must not execute an operator that would trap or that has no meaningful compile-time result. Division by a constant zero, a gather with an out-of-range index, or a reshape whose target shape is inconsistent with the source should raise a diagnostic at compile time rather than produce a poisoned constant. A folder that catches the exception and leaves the node unfolded silently defers the failure to the first request.

Folding also interacts badly with large constants. A fold that materialises a broadcast of a scalar to a full activation shape converts a four-byte constant into a tensor with the same element count as the activation. That is a net loss: the constant now occupies space in the serialized artifact and in device memory, and the broadcast it replaced was free. Compilers therefore bound folding by output size and skip folds whose result exceeds the inputs by a large factor.

Which identities hold in floating point

The second category, algebraic rewriting, replaces a pattern with a cheaper equivalent. The patterns that are always safe are those that do not reassociate or redistribute:

• transpose(transpose(x, p), inverse(p)) is x. The permutation is exact.
• reshape(reshape(x, s1), s2) is reshape(x, s2) when the element counts agree. The mapping is a bijection on indices.
• slice(slice(x, a), b) composes to a single slice whose offsets and lengths are the sum and intersection of the two.
• concat of constants, or of slices that partition the same tensor, reduces to fewer nodes.
• x * 1 is exact for every x, signed zeros and infinities included; the only observable difference is that a signalling NaN is quieted.
• x + 0 is exact except when x is negative zero: under round-to-nearest the sum is positive zero, so the sign of the zero changes.
• x - x is +0 for every finite x, but it is NaN when x is an infinity, so it cannot fold to a zero constant.

The patterns that are not safe are the ones a mathematician would accept without hesitation. (x + y) + z is not x + (y + z): floating-point addition is not associative, and the difference is visible at the last bit and can be amplified by a subsequent reduction. x * (y + z) is not x*y + x*z for the same reason, and distributing a multiply over an add also changes the number of rounding steps. Reassociating a sum of many terms changes the result, which is why a compiler that wants to do it must either be told the change is acceptable or restrict itself to cases where it can prove the change is exact, such as when all operands are integers representable without rounding.

Fusing a multiply and an add into a single fused-multiply-add instruction is the clearest case. FMA computes $x \cdot y + z$ with a single rounding, so it is more accurate than the separate operations, and therefore different. A compiler that contracts a*b + c into an FMA has changed the numerical result of the model in a direction that is usually harmless and occasionally not, because the model was trained with the separate operations. The distinction between “more accurate” and “identical” matters when the comparison is against a reference output.

# Safe: index arithmetic only, no floating-point reassociation.
def simplify_transpose_chain(perm_a, perm_b):
    """Compose two permutations; identity means the pair cancels."""
    if len(perm_a) != len(perm_b):
        return None
    composed = tuple(perm_a[i] for i in perm_b)   # apply perm_a, then perm_b
    return None if composed == tuple(range(len(composed))) else composed

# Unsafe: changes the rounding of a float32 sum.
#   (a + b) + c  ->  a + (b + c)
# Safe only when every operand is an integer with |value| < 2**24,
# where float32 addition is exact and therefore associative.

Rewrite order changes the outcome

Pattern rewriting is not confluent. Two rules can apply to the same node, and applying them in different orders reaches different normal forms. A graph containing reshape(transpose(reshape(x))) may be simplified by collapsing the reshapes first or by pushing the transpose through them first, and one order may expose a cancellation the other does not. Compilers handle this by running a fixed sequence of passes to a local fixpoint and accepting whatever normal form that sequence reaches, rather than searching for a canonical one.

The practical consequence is that pass ordering is part of the compiler’s specification, not an implementation detail. Inserting a new rule early can disable an existing rule by rewriting its pattern into a shape the older rule does not match. This is the reason production compilers organise these rules into named pass pipelines with explicit ordering, as MLIR does with its pass manager and greedy pattern rewriter, rather than as an unordered set of optimisations.

A second ordering hazard is that rewrites can increase graph size temporarily. Pushing a transpose through a chain of elementwise operations multiplies the transpose into each one, which removes the transpose but adds a per-operation permutation attribute. If a later pass then fails to eliminate those, the graph is worse than before. A rewriter with a cost model applies the rewrite only when the estimated cost decreases; a rewriter without one applies it always and relies on a later cleanup pass.

Where the rewrites stop paying

Algebraic simplification has diminishing returns because the remaining nodes are the ones that carry the actual work. Once constants are folded and views are collapsed, what is left is a chain of matrix multiplications, convolutions, normalisations, and activations, and none of those can be replaced by an identity. The pass then contributes almost nothing to the node count, though it has still done its job: it removed the shape plumbing that would otherwise block fusion and inflated the liveness intervals used for memory planning.

There is also a correctness cost that accumulates. Every rewrite is a place where the compiler asserts an equivalence, and the assertion is tested only against whatever reference the project compares outputs with. A rewrite that is exact in the mathematical sense but changes the order of rounding can pass a tolerance check on one model and fail on another whose outputs are near a decision boundary. The conservative position, taken by compilers that must match a reference bit for bit, is to permit only rewrites that are exact in IEEE 754 for all inputs, and to leave the rest behind a flag.

The last thing this pass does is expose constant weights to the next stage. Once a multiply by a constant is known to be a multiply by a specific constant, a compiler can decide to represent it as a scaled integer, to fold it into a preceding normalisation, or to leave it alone. That decision belongs to precision-reduction work rather than to graph rewriting, and the boundary between the two is the point at which the graph stops being a program and starts being a description of specific arithmetic on specific numbers.