Common Subexpression Elimination

Common Subexpression Elimination (CSE) is a compiler optimization that identifies identical expressions computed multiple times within a program's scope, stores the result of the first computation in a temporary variable, and reuses that value for subsequent occurrences. This transforms the computational graph by:

• Replacing redundant nodes with references to a single, cached computation.
• Reducing total FLOPs (floating-point operations) and kernel launch overhead.
• Increasing data locality by reusing values already in registers or cache.

For example, in an expression like y = (a*b) + c; z = (a*b) * d;, the subexpression (a*b) is computed twice. CSE computes it once, stores it as tmp = a*b, and rewrites the code as y = tmp + c; z = tmp * d;.

The power of CSE is defined by its scope—the region of code it analyzes for redundancy.

• Local CSE (Basic Block): Operates within a single basic block—a straight-line sequence of code with no branches. It is fast and simple but misses redundancies across different blocks.
• Global CSE: Analyzes an entire function or procedure, tracking expression values across basic blocks and control flow paths (e.g., loops, conditionals). It uses data-flow analysis to determine if an expression is available (already computed and its value is unchanged) at any given program point. This is more powerful but computationally intensive.
• Interprocedural CSE: Extends analysis across function boundaries, potentially eliminating redundancies in different functions. This is rare in production compilers due to complexity and linking issues.

The standard algorithm for implementing CSE is Hash-Based Value Numbering. It assigns a unique value number to each expression based on its operator and the value numbers of its operands.

• The compiler maintains a hash table mapping expression structures to value numbers.
• When encountering an expression like a + b, it looks up the value numbers for a and b, and the operator +.
• If this (operator, vn_a, vn_b) tuple exists in the table, the expression is redundant and is replaced with the variable holding that value number.
• If not, a new value number is assigned, and the entry is added to the table. This method efficiently identifies structurally identical expressions, even if the original variable names differ (e.g., x+y is identical to a+b if x and a, y and b hold the same values).

CSE is not applied in isolation; it interacts deeply with other compiler passes:

• Constant Folding & Propagation: Creates new opportunities for CSE by replacing variables with constants, making more expressions identical.
• Dead Code Elimination: Must run after CSE to remove the original, now-unused computation if all references to it were replaced.
• Loop-Invariant Code Motion (LICM): Moves invariant expressions outside of loops. CSE can then eliminate redundant instances of those moved expressions.
• Algebraic Simplification: Rewrites x*1 to x or a+a to 2*a, changing the expression's form and creating new CSE opportunities or invalidating previous matches.
• Register Allocation: Benefits from CSE because fewer live variables (the cached results) need to be stored, reducing register pressure.

While powerful, CSE introduces important trade-offs that compilers must balance:

• Increased Register Pressure: Storing temporary results consumes CPU registers. Over-aggressive CSE can cause register spilling to slower memory, negating performance gains.
• Code Size vs. Speed: CSE typically reduces executed instructions but may increase the number of temporary variable assignments, slightly growing code size.
• Side Effects and Aliasing: CSE is unsafe if the redundant expression has side effects (e.g., a function call) or if its operands can be modified between computations (via pointer aliasing). Compilers must perform alias analysis to ensure safety.
• Floating-Point Non-Associativity: For floating-point math, CSE is generally safe and beneficial, as it guarantees bitwise-identical results by reusing the exact computed value, unlike algebraic reassociation.

In machine learning frameworks like TensorFlow, PyTorch (via TorchScript/TorchDynamo), and Apache TVM, CSE is applied to the computational graph (IR) of neural networks.

• Targets High-Level Operators: Identifies redundant transpose, reshape, or slice operations, or repeated constant tensors used in multiple layers.
• Hardware-Aware Decisions: May be disabled for very small, cheap operations where the cost of storing/loading a temporary exceeds the computation cost.
• Fusion with Constant Folding: Often combined into a single pass that folds constants and eliminates their common subexpressions.
• Example: In a vision transformer, the same positional encoding tensor might be added in multiple residual branches. CSE computes it once and reuses it, saving memory bandwidth and compute.

CSE works by analyzing a computational graph's dataflow. It identifies nodes that compute the exact same mathematical expression from the same input tensors. Once identified, the compiler eliminates all but the first computation and rewires the graph so all dependent nodes consume the single cached output. This transforms O(n) redundant work into O(1) work plus a memory read.

• Key Insight: The optimization relies on the deterministic nature of pure computational nodes; the same inputs always produce the same outputs.
• Scope: Can be applied locally within a basic block or globally across an entire function or graph.

In neural network inference, CSE directly targets and reduces wasteful computation, leading to tangible performance gains.

• Latency Reduction: Eliminates duplicate operations, decreasing end-to-end execution time.
• Memory Bandwidth Savings: By reusing a cached tensor, it reduces the number of reads/writes to high-latency memory (e.g., DRAM), which is often the bottleneck.
• Power Efficiency: Fewer arithmetic operations and memory accesses translate directly into lower energy consumption, a critical metric for edge and mobile deployment.

Example: In a vision transformer, the same sinusoidal positional encoding might be computed multiple times for different attention heads. CSE computes it once and shares the result.

CSE is implemented as a pass within modern AI compilers like MLIR, TVM, XLA, and PyTorch's TorchScript/Glow. The pass typically runs after the graph has been canonicalized (put into a standard form) but before hardware-specific lowering.

1. Pattern Matching: The compiler traverses the Intermediate Representation (IR), using pattern matching to find identical expression trees.
2. Value Numbering: A more advanced technique assigns a unique 'value number' to each computed expression; identical numbers indicate redundancy.
3. Graph Rewriting: The IR is rewritten, deleting redundant nodes and redirecting edges.

It often works in concert with passes like constant folding and dead code elimination.

While generally beneficial, applying CSE involves engineering trade-offs that compilers must model.

• Memory vs. Compute: CSE trades compute for memory. The cached result must be stored, potentially increasing the peak memory footprint. The compiler's cost model must decide if the memory pressure is worth the compute savings.
• Fusion Opportunities: An aggressively eliminated subexpression might have been a candidate for operator fusion with a subsequent operation. CSE can sometimes inhibit other optimizations.
• Side Effects: CSE is only safe for pure, side-effect-free operations. It cannot be applied to operations with randomness (e.g., dropout at inference) or stateful reads/writes.

CSE is a member of a family of classical compiler optimizations adapted for computational graphs.

• Constant Folding: A specialized form of CSE where the subexpression consists entirely of constants, allowing the value to be precomputed at compile time.
• Dead Code Elimination: Works synergistically with CSE; after eliminating a redundant calculation, the original node may become dead code if no other consumer exists.
• Common Expression Elimination (CEE): A broader term that may include eliminating similar expressions that can be transformed into a common form, not just identical ones.
• Loop-Invariant Code Motion: A related optimization that hoists identical expressions out of loops, similar to global CSE across loop iterations.

Consider a simple graph snippet computing an element-wise operation: Y = (A + B) * C and Z = (A + B) / D.

Before CSE:

codeCopy
    Add_1 = A + B
    Mul_1 = Add_1 * C
    Add_2 = A + B  // Redundant calculation
    Div_1 = Add_2 / D

After CSE:

codeCopy
    Add_1 = A + B  // Single computation
    Mul_1 = Add_1 * C
    Div_1 = Add_1 / D  // Reuses output of Add_1

The second addition (Add_2) is eliminated. The Div_1 operation is rewired to consume the output of Add_1. This saves one arithmetic operation and the associated tensor allocation.

A compile-time optimization that evaluates and replaces expressions consisting entirely of compile-time constants with their precomputed result.

• Key Mechanism: The compiler performs arithmetic at compile time.
• Example: An operation like y = (2.0 * 3.14) + x is transformed to y = 6.28 + x.
• Impact: Eliminates runtime computation for static values, simplifying the graph and reducing kernel launches.

An optimization pass that identifies and removes operations whose outputs do not affect the final result of the program.

• Key Mechanism: Performs liveness analysis or graph reachability from output nodes.
• Example: Removing a debugging print operation or an unused branch of a conditional.
• Impact: Reduces execution time, memory usage, and binary size by pruning unnecessary computations.

Also known as kernel fusion, this technique combines multiple sequential operations into a single, more efficient kernel.

• Key Mechanism: Fuses ops like Convolution + BatchNorm + ReLU into one compiled kernel.
• Primary Benefit: Drastically reduces intermediate tensor writes to memory (DRAM), which is often the performance bottleneck.
• Use Case: Critical for achieving peak performance on accelerators (GPUs/NPUs) by minimizing memory bandwidth usage.

Identifies redundant calculations of identical expressions and replaces them with a reference to a single computed value.

• Key Mechanism: Creates a hash of operation signatures and input tensors to find duplicates.
• Example: If sin(x) is computed in two separate branches, CSE computes it once and reuses the result.
• Impact: Reduces computational work and can decrease memory footprint by reusing buffer allocations.

A low-level compiler technique that examines small sequences of instructions ('peepholes') and replaces them with more efficient equivalents.

• Key Mechanism: Pattern-matches on a localized segment of the IR.
• Examples:
  • Replacing x * 1 with x.
  • Replacing a sequence of add, mul with a fused multiply-add (FMA) instruction if supported.
• Impact: Produces tighter, faster machine code by leveraging hardware-specific idioms.

A compile-time optimization that pre-allocates and reuses memory buffers for tensors by analyzing their lifetimes in the computational graph.

• Key Mechanism: Performs liveness analysis to determine when tensors are first used and last needed.
• Benefit: Minimizes dynamic memory allocation overhead and reduces the peak memory footprint, which is essential for memory-constrained edge devices.
• Outcome: Creates a static memory plan where buffers are reused for non-overlapping tensors.