As outlined previously, optimizing the computation graph involves transforming its structure to enhance efficiency. These transformations, such as operator fusion or algebraic simplification, are not applied randomly. Instead, they rely on systematic mechanisms known as Graph Rewriting Systems. These systems provide the foundational framework for identifying specific patterns within the graph and applying predefined transformations to replace them with more optimal structures.

At its core, a graph rewriting system operates on the computation graph, typically represented as a Directed Acyclic Graph (DAG). In this DAG, nodes represent operations (like convolutions, matrix multiplications, or activations) and edges represent data dependencies, usually tensors flowing between operations.

Components of a Graph Rewriting System

A graph rewriting system primarily consists of two components:

A set of Rewrite Rules: Each rule defines a specific transformation. It typically has two parts:
- Pattern (or Left-Hand Side - LHS): Describes a subgraph structure to search for. This pattern specifies the types of operations, their connectivity (dependencies), and potentially constraints on node attributes (like data types or tensor shapes).
- Replacement (or Right-Hand Side - RHS): Defines the new subgraph structure that should replace the matched pattern. This might involve creating new nodes, removing existing nodes, and adjusting the edges (data dependencies) to connect the replacement subgraph correctly into the main graph.
An Engine: This component orchestrates the rewriting process. Its responsibilities include:
- Pattern Matching: Efficiently searching the entire computation graph to find all occurrences of the patterns defined in the rewrite rules.
- Rule Application: Selecting which matched rule to apply and performing the graph modification as specified by the rule's replacement pattern.
- Strategy: Determining the order and iteration strategy for applying rules. Applying one rule might enable or disable other rules, so the strategy significantly impacts the final optimized graph.

Pattern Matching

Pattern matching is the process of finding subgraphs within the larger computation graph that are isomorphic (structurally equivalent) to the pattern defined in a rewrite rule. This process needs to consider:

Node Types: Matching the specific operations (e.g., Conv2D, ReLU, Add).
Connectivity: Ensuring the data dependencies between nodes in the subgraph match the pattern's specified edges.
Attributes and Constraints: Verifying that node attributes (like strides, padding, data types, or even constant values) satisfy any constraints specified in the pattern. For example, a fusion pattern might only apply if the data types match or if an operand is a constant.

Matching can range from simple checks for adjacent nodes to complex subgraph isomorphism problems. Techniques often draw from term rewriting and compiler optimization passes, sometimes employing specialized pattern description languages (like MLIR's Pattern Description Language - PDL) to formally define complex patterns and their associated constraints.

Replacement and Transformation

Once a pattern is matched, the rewriting engine applies the transformation defined by the rule's replacement part. This involves careful manipulation of the graph data structure:

Node Creation: New nodes specified in the replacement pattern are instantiated.
Node Deletion: Nodes that were part of the matched pattern but are not part of the replacement are removed.
Edge Rewiring: Edges (data dependencies) are updated. Inputs to the original matched subgraph might be rerouted to inputs of the replacement subgraph. Outputs of the replacement subgraph are connected to the consumers of the original matched subgraph's outputs.

It is essential that the transformation maintains the semantic equivalence of the computation (or improves it, e.g., by removing redundant operations) and the overall validity of the graph structure.

Application Strategies

Since multiple rules might match different parts of the graph simultaneously, or applying one rule might create opportunities for others, the strategy for applying rules is important. Common strategies include:

Iterative Application: Repeatedly scanning the graph and applying rules until no more rules can be applied (reaching a fixed point).
Priority-Based Application: Assigning priorities to rules, ensuring that potentially more impactful or prerequisite transformations are applied first.
Cost Modeling: Using a cost model to estimate the performance impact (e.g., reduced execution time, memory usage) of applying a rule, guiding the selection process towards transformations likely to yield the most benefit.

Greedy application (applying the first matched rule) is often used for simplicity but doesn't guarantee finding the most optimal graph configuration, which is generally an NP-hard problem.

Example: Algebraic Simplification

Consider a simple algebraic identity: Transpose(Transpose(A)) = A. A rewrite rule can capture this:

Pattern: A Transpose operation whose input is the output of another Transpose operation.
Replacement: Replace the outer Transpose operation with a direct connection from the input of the inner Transpose to the consumers of the outer Transpose. Effectively, both Transpose nodes are bypassed and potentially removed if they have no other users.

The rewrite rule identifies two consecutive Transpose operations and replaces them with a direct connection from the original input to the final output.

Example: Simple Fusion

Another common pattern targets fusing an operation with its subsequent activation function, like Convolution -> ReLU.

Pattern: A ReLU operation whose input is the output of a Conv operation. Often includes constraints, e.g., the Conv output is only used by this ReLU.
Replacement: A single FusedConvReLU operation that takes the inputs of the original Conv and produces the output equivalent to the original ReLU. The original Conv and ReLU nodes are removed.

Fusing a Conv and ReLU operation into a single FusedConvReLU node reduces kernel launch overhead and improves data locality.

Graph rewriting systems provide the essential machinery for implementing the advanced graph-level optimizations discussed in the following sections. By defining patterns and their corresponding replacements, compilers can systematically restructure ML computation graphs for improved performance before delving into lower-level code generation. Understanding these systems is fundamental to appreciating how high-level optimizations like operator fusion, algebraic simplification, and layout transformation are realized in practice.