Constant Folding

Constant folding evaluates arithmetic and logical expressions composed entirely of compile-time constants, replacing them with a single constant tensor. This eliminates runtime computation for that subgraph. For example, the expression (weight * 0) + bias where weight and bias are known constants would be folded directly into the resulting bias tensor, removing the redundant multiplication and addition operations.

By replacing complex subgraphs with precomputed constants, folding reduces the number of nodes in the computational graph. This leads to:

• Fewer kernel launches and reduced scheduling overhead.
• Lower memory bandwidth usage, as intermediate tensors are eliminated.
• A simpler graph that enables more aggressive downstream optimizations, such as dead code elimination and better fusion opportunities.

Constant folding is rarely applied in isolation. It works synergistically with other compiler passes:

• Dead Code Elimination: Removes nodes whose outputs are no longer needed after folding.
• Common Subexpression Elimination (CSE): Can create new constant expressions that folding then evaluates.
• Shape Inference: Often a prerequisite, as folding requires knowing the precise shapes of constant tensors.
• Quantization Folding: A specialized form that folds fake quantization nodes with constant weights.

A key characteristic is its strictly compile-time nature. The optimizer must have complete knowledge of all input values to an expression. This differs from optimizations like constant propagation, which may handle some runtime values. Folding is most powerful in Ahead-of-Time (AOT) compilation scenarios for deployment, where model weights and many hyperparameters are fixed and known.

Common patterns in models where constant folding applies:

• Fixed normalization layers: (input - mean) / std where mean and std are fixed.
• Static rescaling ops: Multiplying by a fixed scalar after a layer.
• Constant padding with a known value.
• Fixed affine transformations in positional encodings or embeddings.
• Compile-time shape manipulations like reshape, slice, or tile with constant arguments.

Constant folding is not always applicable or beneficial:

• Dynamic Inputs: Cannot fold expressions dependent on runtime input data.
• Control Flow: Operations inside conditionals or loops with constant conditions can sometimes be folded, but analysis is complex.
• Memory vs. Compute Trade-off: Folding a large tensor operation (e.g., a big matrix multiply) into a single constant can increase the model's static binary size, trading compute for storage. The optimizer's cost model must evaluate this balance.
• Numerical Stability: Must ensure the compile-time evaluation uses the same precision and semantics as the runtime would.

A primary application is simplifying arithmetic operations on static model parameters. For example, a graph containing a sequence like (Weight * 2.0) + 1.5 is evaluated at compile-time. The entire expression is replaced by a single, precomputed constant tensor, removing the multiplication and addition operations from the runtime execution path. This is pervasive in models where weights are fixed post-training.

Constant folding optimizes layer normalization and batch normalization during deployment. These layers often involve operations on learned parameters (gamma, beta) and running statistics (mean, variance).

• The affine transformation (x - mean) / sqrt(variance + epsilon) * gamma + beta is composed entirely of constants post-training.
• The compiler folds this into a single, efficient scale-and-shift operation: x * folded_scale + folded_bias, eliminating multiple arithmetic kernels.

Conditional logic based on immutable configuration tensors is resolved at compile-time. For instance, a graph may contain a branch like if (training_mode_tensor == False): path_A else: path_B. Since training_mode_tensor is a constant False during inference, the compiler:

• Folds the conditional check.
• Prunes the entire path_B subgraph as dead code.
• Retains only path_A, streamlining the inference graph. This is common in frameworks that unify training and inference graphs.

In models with static input embeddings, constant folding can cache the results of initial embedding layers. If the first network stage is a lookup from a fixed embedding table followed by a constant operation (e.g., a fixed positional encoding addition), the compiler can precompute the combined result for all valid input indices. This transforms an initial dynamic lookup into a simple, folded constant fetch, reducing latency for the first network layer.

Sequences of static shape transformations are folded into a single, direct mapping. Consider a tensor flow: reshape(permute(reshape(constant_tensor, shape_A), axes), shape_B). If constant_tensor, shape_A, shape_B, and axes are all compile-time constants, the compiler:

• Calculates the final output tensor directly in memory.
• Replaces the entire operation chain with the final, pre-laid-out constant.
• This eliminates the memory bandwidth overhead of intermediate rearrangements.

Constant folding is not applied manually but is a core pass in ML compilers like Apache TVM, Google's XLA, and Intel's OpenVINO. Its impact is measured by:

• Reduced Operation Count (OPs): Directly decreases the number of kernels launched.
• Simplified Graph Topology: Leads to cleaner graphs, enabling further optimizations like operator fusion.
• Deterministic Latency: By moving computation to compile-time, it removes variable arithmetic latency from the critical inference path, improving predictability.