Equality Saturation

The IRs we discussed above are all used to analyze and transform the underlying program structure, so as to produce a new representation of the optimized program. In a con- ventional optimizing compiler, program optimization is often carried out in a sequence of transformation passes, where each pass accepts a program, often written in a certain IR, and produces an optimized program in the IR. The traditional practice is to always apply these optimization phases in a fixed order, but a good ordering of these phases is crucial to achieve a good optimized result, and the optimal ordering varies across applications being compiled [Alm+04]. The process of finding the optimal ordering is known as the phase-ordering problem, which is in general undecidable [TB06]. Moreover, programs running on CPUs or GPUs are usually quantified by their throughput or latency, in contrast, designs on FPGAs concern us with additional objectives besides run time, such as power consumption and resource utilization that impact the quality of synthesized circuit. Multiple designs which trade-off these objectives could exist, and which design to choose relies on the specifics of the use case. It is therefore sensible to explore the design space by optimizing multiple objectives simultaneously. For the above reasons, it is desirable to have an IR and the associated optimization procedures to efficiently discover equivalent structures that lead to different implementations of the original program.

In software, a novel approach called equality saturation is proposed in [Tat+] to find multiple possible optimized variants of the original program, and subsequently deal with the phase-ordering problem. It creates a new graph-based IR, program expression graph (PEG), to encode the effects of executing the program.

To begin, we review the structure of the PEG, by considering a simple loop example in Figure 2.11. By understanding how the PEG can be evaluated for the output values, we can interpret how the PEG captures the control- and data-flow information of the program. PEGs, similar to arithmetic expressions expressed in a tree structure, can be evaluated in a bottom-up fashion, by recursively propagating computed values from the leaf nodes to the root of the tree. However, unlike arithmetic expressions which are acyclic, edges in PEGs may form cycles to express loops in the original program.

int x = 1; int y = 1; while (y <= 10) { x = y * x; y = y + 1; }

(a)The original program.

x eval1 θ1 1 _* y eval1 θ1 1 + 1 pass₁ ¬ ≤ 10 (b)The resulting PEG.

Figure 2.11. A simple loop which computes the factorial of 10, and the resulting PEG. This example and its PEG, showing computations that lead to the final x and y, is taken from [Tat+].

Data-Flow Nodes

All loops in the PEG are formed by θ nodes, which is used in the following form:

θ

ι func( )

where it accepts two child graphs, ι and func, and func further takes the θ node as one of its inputs to form a complete cycle. Evaluating a θ node produces a list of values computed iteratively by the node’s subgraph. The first value in the list, is the computed result of ι, which we name i, and values in the rest of the sequence are iteratively computed by func. In functional programming, this is similar to iteratively computing the fixpoint fix F of an initial empty list [], which is defined as:

fix F ([]) = lim

n→∞F

Here, map(func, v) applies the subgraph computation func to all elements in the list v, and prepend(i, v0)prepends the element i to the list v0.

For example, the following subgraph extracted from Figure 2.11b evaluates to the sequence [1, 2, 3, 4, . . .]:

θ1

1 +

1

It is noteworthy that θ nodes may have subscripts. For instance, in Figure 2.11b, both nodes θ1 share the same subscript 1. This is used to indicate that the two sequences produced by both θ nodes iterate simultaneously, i.e. they share the same iteration count so that a new value for x can be computed as we update y. Therefore, the θ1 node in the left of this figure produces a sequence of the factorials of [1, 2, 3, . . .], i.e. [1, 2, 6, 24, . . .].

Computation nodes, such as arithmetic + and Boolean operators ≤ and ¬ in Figure 2.11b, operates on a list of values, by performing the computation on each value in the list.

For instance, the ≤ node accepts two inputs, the sequence [1, 2, 3, . . .] derived earlier, and a scalar 10, computes the result of x ≤ 10 for each element x in the sequence, and finally produces the list, where tt and ff respectively denote true and false Boolean values:

[tt, tt, tt, tt, tt, tt, tt, tt, tt, tt, ff , ff , ff , . . .]. (2.53)

The subsequent ¬ node then negates all elements in the list:

[ff , ff , ff , ff , ff , ff , ff , ff , ff , ff , tt, tt, tt, . . .]. (2.54)

Control-Flow Nodes

The θ node encodes an infinite sequence of computed values, whereas the output value of the program is a scalar. By further representing control-flow information in a PEG, it

becomes possible to refer to a single value in this sequence, as the output of the program. To do so, Tate et al. [Tat+] further introduce pass and eval nodes. The pass node finds the first true (tt) value in a sequence of Boolean values, and returns the index of this value, and eval takes two child nodes, where the first node evaluates to a listv of values, and the second is an scalar n used to select a scalar value from l, as the output of this node.

To illustrate, the pass₁ node finds the first tt in the list (2.54), 11. The eval node of the output variable y subsequently fetches the 11-th item from the list [1, 2, 3, 4, . . .] we produced earlier, which is 11. Similarly we can apply the same process to find that the output x is 10 ! , i.e. the factorial of 10.

By using pass and eval nodes to represent the control-flow in an algebraic fashion, and mixing data- and control-flows together in the graphical representation, PEGs provide us with greater opportunities to optimize data-flow across control-flow boundaries andvice versa. Simple equivalence rules can be defined for these nodes algebraically. For instance, arithmetic operators can be distributed over θ and eval, e.g. eval(j, k) + i ≡ eval(j + i, k). Complex transformations can therefore be deductively constructed from these simple equivalence rules.

Equivalence Finding

By applying transformation passes to the PEG, their approach detects incremental modifi- cations, and appends these changes to the original PEG. The new changes, represented as extra structures in the PEG, are linked to their corresponding equivalent nodes by equivalence edges. These edges indicate a pair of subgraphs are equivalent, forming a program expression graph with equivalent structures (E-PEG), that could capture multiple PEGs in the same structure. The resulting E-PEG is similar to the one in Figure 2.12, where dashed edges indicate equivalences. It is notable that each edge allows a binary choice, therefore the number of PEGs that can be represented in an E-PEG could be exponential in the number of equivalence edges.

∗ θ 5 0 φ δ + 3 + 1 θ ∗ 0 0 5 ∗ 5 φ δ _∗ 5 + ∗ 15 ∗ 3 5 5 + ∗ 5 1 5

Figure 2.12. An simple E-PEG example, taken from [Tat+].

By repeatedly applying transformation rules to append all possible equivalent structures to the E-PEG, this graph will eventually saturate, i.e. no more equivalent structure can be added to the graph because all possible equivalences are now discovered. This process and the resulting E-PEG is more space-time efficient than enumerating all possible PEGs along the path, because E-PEG encourages sharing common subgraphs, even across equivalent edges. This saturated graph can always be produced regardless of in what order we apply the passes, hence preventing the phase-ordering problem. Furthermore, E-PEG defers the decision of whether an optimization should be committed until we have reached full saturation, allowing the global optima to be discovered. In contrast, because each optimization pass in a conventional compiler is performed consecutively, the compilers must make the decision to commit changes immediately after applying the optimization, which consequently often results in local optima.