Canonicalization

The primary goal is to remove duplicate or unnecessary computations from the graph. This includes:

• Algebraic simplification: Applying mathematical identities (e.g., x * 1 = x, x + 0 = x).
• Constant folding: Pre-computing operations on constant tensors.
• Common subexpression elimination (CSE): Identifying and caching identical subgraphs.

For example, an expression like (a + b) * 1 + 0 is canonicalized to simply a + b, removing the redundant multiplicative identity and additive zero operations.

Canonicalization enforces a single, consistent representation for semantically equivalent operations. This is critical for pattern matching in later passes like operator fusion or hardware delegate selection.

Examples include:

• Converting all convolution padding specifications to a single format.
• Ensuring all reduction operations (sum, mean) use the same axis parameter representation.
• Rewriting a transpose followed by a matrix multiplication as a single batched matmul with permuted indices.

Without this standardization, the compiler would need separate pattern rules for every syntactic variation of the same operation.

This objective transforms complex, nested conditional and loop structures into a normalized form. It involves:

• Flattening nested conditionals into a sequence of simpler if-else blocks.
• Canonicalizing loop bounds to use zero-based indexing and unit strides where possible.
• Removing dead branches identified through static analysis and constant propagation.

This simplification is a prerequisite for advanced loop optimizations like loop tiling, unrolling, and vectorization, which require predictable, regular control flow patterns.

By producing a predictable, canonical graph from potentially diverse input representations, this pass ensures that downstream optimizations are applied consistently and reliably. This is essential for:

• Reproducible builds: The same model always compiles to the same efficient executable.
• Regression testing: Performance changes can be traced to specific optimization modifications, not graph representation quirks.
• Cross-framework compatibility: Models imported from PyTorch, TensorFlow, or ONNX are reduced to a common internal representation before hardware-specific optimizations.

Canonicalization actively reduces the number of nodes and edges in the computational graph, lowering the search space for subsequent NP-hard optimization problems like graph partitioning or optimal operator scheduling.

Techniques include:

• Fusing chains of element-wise operations (e.g., ReLU -> Add -> Clip) into a single composite node.
• Eliminating identity operations like a Reshape that does not change the tensor layout.
• Propagating shape information via shape inference to remove dynamic guards.

A simpler graph allows the compiler's cost model to evaluate optimization choices more quickly and accurately.

This final objective transforms the graph into a form that is ideal for the graph lowering process to hardware-specific instructions. It creates a clean separation between high-level, framework-agnostic optimizations and low-level, backend-specific code generation.

Key preparations include:

• Type normalization: Ensuring all tensors have explicit, concrete data types (e.g., float32, int8).
• Explicit broadcasting: Making implicit broadcast operations explicit with Expand or Tile nodes.
• Layout annotation: Marking tensors with preferred memory layouts (e.g., NHWC or NCHW) for data layout optimization.

This creates a well-defined contract between the canonical intermediate representation (IR) and the backend compilers for GPUs, NPUs, or CPUs.

Applies mathematical identities to simplify expressions, reducing computational cost. This includes:

• Constant folding: Evaluating operations with constant inputs (e.g., Add(5, 3) → 8).
• Identity elimination: Removing operations that have no effect (e.g., Multiply(x, 1) → x).
• Zero propagation: Simplifying operations involving zeros (e.g., Add(x, 0) → x).
• Strength reduction: Replacing expensive operations with cheaper equivalents (e.g., Pow(x, 2) → Multiply(x, x)).

Identifies and eliminates redundant calculations of identical expressions within the graph. When the same computation appears multiple times with the same inputs, CSE:

• Caches the result of the first computation.
• Replaces all subsequent identical nodes with a reference to the cached result.
• Crucially reduces compute and memory overhead in graphs with repeated patterns, such as those generated from high-level frameworks or after unrolling loops.

Removes operations whose outputs do not contribute to the model's final output. The compiler performs a live variable analysis to trace data dependencies from output nodes backward. Any operation not on a path to a required output is deemed dead code and pruned. This eliminates:

• Unused branches or layers from experimental architectures.
• Debugging or profiling operations not needed for inference.
• Intermediate calculations made redundant by other optimizations.

Enforces a consistent, canonical form for all operations in the graph. This involves:

• Operator canonicalization: Converting different operator representations that perform the same function into a single, standard form (e.g., normalizing all convolution padding specifications).
• Attribute normalization: Ensuring operator attributes (like axis parameters) use a consistent indexing scheme (e.g., always 0-based).
• Tensor layout canonicalization: Deciding on a default data format (e.g., NHWC or NCHW) and inserting explicit transpose operations where necessary, making subsequent data layout optimizations predictable.

Transforms high-level, structured control flow into a linearized, primitive representation. This is essential because many hardware backends and low-level IRs operate on basic blocks. The pass:

• Lowers complex control constructs like tf.while_loop or torch.jit.script conditional branches into primitive Switch, Merge, and NextIteration nodes.
• Flattens nested scopes or subgraph structures into a single, flat list of nodes.
• Explicitly materializes iteration counts and loop conditions, making dataflow and memory planning analyzable by subsequent passes.

A foundational pass that annotates every node in the graph with the data type and dimensional shape of its output tensor. This is not just inference; it's a canonicalization step that:

• Resolves polymorphic operators into their concrete, instantiated form based on input types.
• Inserts explicit cast operations where type mismatches occur, eliminating implicit casting semantics of frontend frameworks.
• Propagates static shape information as far as possible, which is critical for static memory planning and enabling aggressive optimizations like constant folding.

A compile-time optimization that evaluates and replaces subgraphs composed entirely of constant values with a single constant tensor containing the precomputed result. This eliminates runtime computation, simplifies the graph, and often enables further optimizations.

• Key Benefit: Removes unnecessary compute nodes, reducing graph size and inference latency.
• Example: An expression like Add(Constant(5), Constant(3)) is replaced with Constant(8).
• Relation to Canonicalization: Often performed after canonicalization, as a standardized graph makes constant sub-expressions easier to identify and fold.

An optimization that identifies and eliminates redundant calculations of identical expressions within a graph. It caches the result of the first computation and reuses it, rather than recalculating.

• Key Benefit: Reduces computational workload and memory operations for repeated patterns.
• Example: If the same normalization is applied to two different tensor branches, CSE will compute it once and share the result.
• Relation to Canonicalization: Canonicalization, by creating a standard form, makes it significantly easier for the compiler to recognize structurally identical sub-expressions for elimination.

A compiler pass that identifies and removes operations whose outputs do not contribute to the final output of the computational graph. This includes unused branches, orphaned constants, and side-effect-free operations.

• Key Benefit: Prunes the execution graph, reducing both memory footprint and inference time.
• Prerequisite: Requires precise dataflow analysis to determine which nodes are 'live'.
• Relation to Canonicalization: A canonicalized graph, with simplified control flow and standardized operators, allows for more accurate and aggressive dead code analysis.

A critical optimization that combines multiple sequential operations (e.g., Convolution, BatchNorm, Activation) into a single, compound kernel. This reduces intermediate tensor writes to memory (kernel launch overhead) and improves data locality.

• Key Benefit: Dramatically reduces memory bandwidth pressure, which is often the bottleneck for neural network inference.
• Hardware-Aware: Fusion patterns are often tailored to specific accelerator capabilities (e.g., GPU tensor cores, NPU instruction sets).
• Relation to Canonicalization: Canonicalization simplifies the graph and exposes clear, sequential patterns of operations that are ideal candidates for fusion.

The process of transforming a high-level, hardware-agnostic intermediate representation (IR) of a model into a lower-level, target-specific IR or machine code. This involves mapping abstract operators to concrete hardware instructions.

• Key Steps: Includes operator decomposition, legalization (converting unsupported ops), and scheduling.
• Target-Specific: The final lowered graph is optimized for a specific backend like CPU, GPU, or NPU.
• Relation to Canonicalization: Canonicalization typically occurs early in the lowering pipeline, creating a clean, predictable high-level graph that subsequent target-specific lowering passes can process reliably.

A low-level optimization technique that examines small, local sequences of instructions or operations (a 'peephole') and replaces them with a more efficient sequence that produces the same result.

• Scope: Works on a very localized scale, often on the instruction or simple operator level.
• Examples: Replacing Add(x, 0) with x, or fusing a Transpose immediately followed by another Transpose.
• Relation to Canonicalization: Canonicalization creates the standardized, simplified patterns that peephole optimizers are designed to recognize and replace, making them more effective.