Yichen Yang, Phitchaya Mangpo Phothilimthana, Yisu Remy Wang, Max Willsey, Sudip Roy, Jacques Pienaar, "Equality Saturation for Tensor Graph Superoptimization", MLSys 2021
Graph Rewriting and Phase-Ordering Problem
Deep learning frameworks often employ ML compilers to transform a given tensor computation graph into a more efficient, semantically equivalent graph. This process, dubbed as graph rewriting, uses a set of curated rewrite rules that matches a certain pattern in a tensor graph and transforms it into an equivalent one. For example, when applied to a given expression a * 2 + 1, a rewrite rule a * 2 -> a << 1 matches the subexpression a * 2 and transforms it into a << 1, resulting in an equivalent expression of (a << 1) + 1.
However, applying these rewriting rules may enable or disable the application of other rewriting rules. For example, if both a * 2 -> a << 1 and (a * 2) / 2 -> a are both present in the set of rewriting rules, (a * 2) / 2 cannot be simplified if the former rule is applied first, substituting the expression into (a << 1) / 2. Thus, the optimization results depend heavily on the order of rewriting rules applied, posing a classic problem known as phase-ordering or rewriting-ordering problem.
Many solutions to the problem have been proposed in the compiler community. This includes:
- Heuristic solutions, which simply predefines the order of applying rewrite rules.
- Search-based solutions, which searches for sequences of substitutions to outperform heuristic approach.
Nevertheless, these two solutions both rely on sequential order of application, and thus consider only a small portion of exponential-sized ordering space. Therefore, a relatively novel approach called equality saturation, which we will elaborate in the next section, has been more widely adopted, thanks to its ability to consider all possible sequences of rewriting.
Graph Saturation
Equality saturation makes use of a data structure called e-graph to efficiently track all possible equivalent graphs. The process consists of two phases, namely, exploration and extraction. In the exploration phase, the given tensor graph expressed in the form of e-graph is iteratively applied with the given rewrite rules to explore all possible ways of rewriting the input program. The extraction phase takes these equivalent expressions and chooses the one with the lowest cost according to the cost model.
E-graph
An e-graph is a set of equivalence classes (e-classes), each of which is a set of equivalent e-nodes.
Take the above e-graph that represents the two equivalent expressions of (a * 2) / 2 and (a << 1) / 2. In the figure, each e-node corresponds to the expression represented by the subtree with itself being the root node.
As in the figure, each dotted box represents an e-class, where each of the element e-nodes represent the semantically same subgraph. For example, since * and << are in the same e-class, they represent the equivalent expressions of a * 2 and a << 1.
More formally, an e-graph is defined as follows:
E-graph
- An e-graph represents a term if any of its e-classes do.
- An e-class represents a term if any of its equivalent e-nodes do. All terms represented by an e-class are
equivalent.
- An e-node $f(c_1,\cdots,c_n)$ represents a term $f(t_1,\cdots,t_n)$ if each e-class $c_i$ represents $t_i$. A childless e-node a represents a constant term $a$.
Applying a rewrite rule to an e-graph yields a "supergraph" that consists of all the e-nodes in the original e-graph, plus the newly added e-nodes from the rewrite. Importantly, this process is non-destructive. Thus, rather than replacing the matched subexpression, the rewritten form is simply added to the same e-class as the original. This allows the e-graph to represent both the original and rewritten expressions at the same time, without losing any information.
1. Exploration Phase
The exploration phase begins by initializing an e-graph from the given tensor graph. Next, all the rewrite rules are applied repeatedly until saturation is reached, i.e., no rewrite rule adds a new information to the e-graph. In practice, the iteration will also be terminated if a user-specified limit is reached, e.g., maximum iteration or maximum e-graph size.
The process enables, ideally, every possible rewriting of the input program to be represented in the e-graph, from which the lowest cost representation will be extracted in the following phase.
Once the exploration phase is complete, the e-graph encodes an exponentially large set of equivalent expressions. The extraction phase selects the best one according to a user-provided cost model, which assigns a cost to each e-node.
Various methods can be employed for the extraction phase, including greedy solutions and integer linear programming (ILP)-based solutions.
Method
Challenges in Applying Equality Saturation to Tensor Graphs
Despite its effectiveness, the authors argue that the integration of equality saturation method into tensor graph rewriting problem is non-trivial for two reasons:
First, existing libraries lack the support for multi-pattern rewrite rules, which is often very effective in optimizing tensor graph optimization in GPU. For example, whenever the two patterns (1) $y_1 = Ax_1$ and (2) $y_2 = Ax_2$ are matched simultaneously for different inputs $x_1$ and $x_2$, one can apply a rewrite rule $(y_1; y_2) = A(x_1; x_2)$ to accelerate the computation by first concatenating the inputs before multiplying the matrix $A$ and splitting the resultant product.
Second, tensor computation graphs are DAGs rather than trees, as opposed to in the classic targets of equality saturation. Therefore existing methods do not account for cycles that may arise in the e-graph during the exploration phase.
Solution: TenSat
To cope with these two problems, the authors devise a novel equality saturation method applicable to tensor graphs, namely, TenSat.
First, TenSat effectively handles multi-pattern rewrite rules by finding matches for each of the required patterns individually, taking Cartesian product of the sets of matches, and verifying if the combination of matches is compatible. Here, the multiple patterns are canonicalized so that we can reduce the overhead for matching. For example, the two required patterns in the previous example, $y_1 = Ax_1$ and $y_2 = Ax_2$, are both canonicalized into $y = Ax$, so that the solver needs to search for matches of this pattern only once rather than separately for each sub-pattern.
Second, to cope with the cycles, the extraction phase employs an additional constraint to make sure the final solution does not have any cycles. Alternatively, the authors also suggest a separate cycle filtering step that can be integrated in the exploration phase, enforcing none of the expressions in the final e-graph contain cycles.
The Algorithm
The exploration phase of the suggested TenSat algorithm is demonstrated in Algorithm 1.
As described previously, the algorithm first canonicalizes all the multi-pattern rewrite rules so that it can reduce the number of pattern matchings performed (line 2-8). Next, it follows the standard exploration phase of equality saturation method, while incorporating additional steps to decanonicalize the matches and check if the combination of each pattern match is compatible. If the combination turns out to be compatible, the rewriting is applied to the e-graph, adding new e-nodes and expanding each e-class in the graph. The iteration proceeds up to maximum of MAX_ITER times.
The extraction phase then selects the optimal graph from the fully explored e-graph. TenSat formulates this as an Integer Linear Program (ILP), where a binary variable $x_i \in {0, 1}$ is introduced for each e-node $i$ to indicate whether it is selected, and the objective is to minimize the total cost $f(x) = \sum_i c_i x_i$. The following constraints are imposed:
- $\sum_{i \in e_0} x_i = 1$: exactly one e-node is chosen from the root e-class.
- $\forall i, \forall m \in h_i, x_i \leq \sum_{j \in e_m} x_j$: if an e-node is selected, at least one e-node in each of its children e-classes must also be selected.
- $\forall i, \forall m \in h_i, t_{g(i)} - t_m - \epsilon + A(1 - x_i) \geq 0$ and $\forall m, 0 \leq t_m \leq 1$: the topological ordering variables $t_m$ enforce that the extracted graph contains no cycles. Here, $\epsilon < 1/M$ is a constant to enforce a stricter inequality.
The optimal solution to this ILP is guaranteed to yield a valid, cycle-free graph with the minimum cost under the given cost model. However, the cycle constraints are known to significantly slow down the ILP solver as the e-graph grows larger. Therefore, when cycle filtering is applied during the exploration phase, the last two constraints can be dropped entirely, leading to substantially faster extraction.
Evaluation
The authors compare the run-time speedup of the optimized graph and the optimization time of TenSat with TASO, a search-based baseline.
As demonstrated in the figure, TenSat showed up to 16% runtime speedup over TASO while boasting optimization time order of magnitudes faster.