Automatic Datapath Optimization using E-Graphs

A useful example of transformation-based datapath improvement comes from Verma, Brisk and Ienne [4], which automatically applies data-flow transformations to maximally exploit carry-save representation. Their primary objective is to cluster additions together in the data-flow graph, a useful target as full carry-propagate addition is generally expensive and can often be avoided. This can be done by deploying compressor trees, circuits taking three or more input words which get reduced to two output words: a carry and a save. Using a carry-propagate adder to sum the carry and the save returns the sum of all the inputs. This can be beneficial as, for example, combining two consecutive carry-propagate adders into a compressor tree and one carry-propagate adder saves the cost of one carry-propagate adder at the expense of a compressor tree. We generalize this work using our methodology, which is able to replicate the results obtained in [4] as a special case.

Another well-studied transformation beyond the reach of standard commercial synthesis tools is the multiple constant multiplication (MCM) problem [6, 7]. The MCM problem asks, given a set of integer coefficients $\{a_{1},...,a_{n}\}$ and variable $x$, what is the optimal architecture to compute the set $\{a_{1}\times x,...,a_{n}\times x\}$? Competing solutions use a fixed number representation of the constants [7], often canonical signed digit (CSD) representation [8], and/or deploy an adder graph algorithm [6]. A transformation based approach also captures both of these methods.

In addition to these special cases, there is a wide variety of transformations that can be captured through standard arithmetic rewrites, e.g. associativity, distributivity, etc. Often these rewrites interact with each other, in the sense that applying one type of transformation opens or closes the door to applying a different class of transformation.

Equivalence graphs, commonly called e-graphs, provide a dense representation of equivalence classes (e-classes) of expressions [9]. Often found in theorem provers, this data structure enables a graph optimization technique called equality saturation [5, 10, 11]. The e-graph represents expressions, where the nodes, known as e-nodes, represent functions (including variables and constants, as 0-arity functions) and are partitioned into a set of e-classes. The intuition is that e-classes can be used to compactly represent equivalent expressions, whose evaluation always leads to the same result. Edges represent function inputs and are from e-nodes to e-classes; see Figure 1, where dashed lines represent e-class boundaries, solid ellipses are e-nodes and arrows are edges. We define $\mathcal{C}$ to be the set of e-classes, $\mathcal{N}$ the set of e-nodes and $E\subseteq\mathcal{N}\times\mathcal{C}$ the set of edges. We also introduce $\mathcal{N}_{c}$ to denote the set of e-nodes in a given e-class $c$.

Rewrites define equivalences over expressions, for example $x+x\rightarrow 2\times x$ says that $x+x$ is equivalent to $2\times x$. Such rewrites are applied constructively to the e-graph, meaning that the left hand side of the rewrite remains in the data structure. Constructive rewrite application avoids the concern of which order to apply rewrites in. As rewrites are applied, the e-graph grows monotonically, representing more and more equivalent expressions, and hence naturally capturing the interaction between different rewrite rules.

Equality saturation provides us with a stopping condition. At the point where further rewrites add no additional information, we say that the e-graph has saturated. From an e-graph representing potentially infinitely many equivalent expressions we may choose the “best” expression [5].

egg is a recent Rust e-graph library, which is intended to be a general purpose and reusable implementation [5]. It adds powerful performance optimizations over existing, usually bespoke, e-graph implementations along with some useful additional features. It has been used to automatically improve the numerical stability of floating point expressions [12], map programs onto hardware accelerators [13] and optimize linear algebra [14]. To build a functioning e-graph optimization tool, egg must be supplied with a language definition – that is a set of operator names together with their arity, and a rewrite set – that is a set of equivalences over the given language definition.

We consider the problem from the abstraction level of a Verilog parser and operate on finite length bitvectors. We target bitvector arithmetic and bitwise Boolean operations as they form the basis of low-level datapath optimization. Including the bitwidth of these vectors is crucial to correctly model the circuit’s behaviour. Bitwidths are also essential to correctly evaluate the cost of a given operation, clearly an 8-bit addition is less expensive than a 32-bit addition. The defined operations represented by nodes in the e-graph are described in Table I, and are defined for all inputs.

The first set in the table is a subset of the operators defined in Verilog [19]. These operators are fundamental in most arithmetic circuit designs and manipulating designs using them can have significant effects on power, performance and area. Bitwidth information is included in the e-graph as edge labels between the nodes. In addition to these basic operators, we introduce a second set of operators in Table I (below the dashed line), which encode the merging and sharing capabilities of modern synthesis tools. In [4], the benefits of merging adjacent additions to maximise the use of carry-save are exploited. We are able to capture the same effect in an automated manner through introduction of a SUM node in our language definition, which combines an arbitrary number of additions into a single node. This node encodes the compressor tree and carry-propagate adder described in Section I. Figure 3 shows how consecutive additions can be rewritten as a SUM node. Other such merging operators are the fused multiply-add (FMA) and the less familiar muxed multiplication array (MUXAR), which blends two disjoint multiplication arrays into one. These are discussed further in Section III-B.

We have identified and formalized rules that capture many of the manual transformations that Intel’s Numerical Hardware Group regularly deploy manually. The set of rewrites described in Table II define equivalences over expressions. Each of the rewrites incorporates bitwidth information, where bitwidths are associated with operands. We introduce a left subscript notation, $\prescript{}{p}{x}$, to denote a bitvector $x$ with length $p$ bits. Left subscript takes the lowest operator precedence.

egg supports conditional rewrites. We make use of this facility to allow us to decide whether a rule may be applied based on its bitwidth information. The sufficient (but not always necessary) conditions under which each rule applies are described in Table II. As an example, consider the “Distribute Mult Over Add” rewrite, where the sufficient condition is a function of the bitwidth parameters, $\min(q,u,v)\geq r$. When this condition is satisfied the rewrite can be safely applied as the equivalence holds. For some rewrites, a parameter appears on the right hand side but not on the left, for example $q$ in the associativity rewrites. This corresponds to a condition constraining but not fully determining the values that these undefined parameters can take, in which case we select the minimum feasible bitwidth. The conditions have been validated by creating a parameterizable SMT query for each rewrite using the theory of fixed size bitvectors [20], which have been checked for all combinations of bitwidths in {1,…,10}, however a formal equivalence check between the generated and original RTL ensures that we do not need to trust the correctness of our rewrites or of the egg library.

Arithmetic Logic Exchange rewrites are inspired by a set taken from [4], which focus on the interplay of logic and addition. To extend the prior work we have added rewrites that are able to move other arithmetic operators and we have generalized the rules so they are able to be applied in a multiple bitwidth setting.

The “Merging Ops” rewrites allow us to correctly evaluate the cost of specific sequences of operations, that logic synthesis tools are able to effectively optimize. We have seen that the SUM node merges consecutive additions, highlighting that the area usage for consecutive additions is not additive. This captures what logic synthesis tools will do to an expression such as $a+b+c$, converting it to a compressor tree and deploying a single carry-propagate adder. Further merging optimization capabilities of Synopsys Design Compiler are documented [1]. The “Merge Mult Array” rewrite does not make use of carry-save format but identifies that two disjoint multiplication arrays can be merged. Letting $a[i]$ represent bit $i$ of $a$ and $u=\lceil\log_{2}(r)\rceil$, MUXAR is shorthand for the right hand side of the rewrite, where the SUM represents array reduction:

Logic synthesis is capable of exploiting this optimization if it identifies an expression of the form $(a\times b)+(c\times\sim b)$, but we must indicate the merging opportunity to our tool.

We added the “Constant Expansion” rules to re-express multiplication of a variable by a constant. These rules allow us to recreate results from the literature described in Section II-A on the MCM problem. In addition to the explicit rewrite rules, constant folding is implemented as an e-class analysis in egg [5].

Having applied rewrites to the e-graph until saturation or timeout limits are reached, we now must extract the optimal design from potentially infinitely many choices. The extraction process must select a set of e-classes to implement and for each e-class, which e-node within that class to implement, subject to the constraint that the e-class children of each selected e-node must also be selected. Choosing an optimal design requires some metric that allows us to discriminate between competing implementations. Industrial circuit design is typically judged on area, latency and power consumption. In this contribution we will only use an area metric, therefore our definition of optimal will be the smallest circuit implementation.

We have developed a theoretical area estimate cost function in terms of the number of two-input gates required for the operator. It assigns a cost per operator that is a function of the input and output bitwidths. Table I lists the architectures on which we base the cost of the more complex operators. We introduce different costs for when at least one of the operands is a constant. The cost of a complete design is then the sum of the operator costs, and the objective is to minimise this cost.

The major benefit of using a theoretical cost metric as opposed to using metrics derived from logic synthesis or HLS tools is the computation speed which, when combined with equality saturation, enables effective design space exploration. Cost metric validation is addressed in Section V.

By explicitly introducing rules for operator merging and sharing, we are able to define a cost for each node based only on its type and argument bitwidths, capturing downstream synthesis optimizations for SUM, MUXAR and FMA.

When extracting an RTL implementation from an e-graph, one is immediately faced with the question of common subexpressions. Common subexpressions are naturally extracted as part of the e-graph construction process [5] and ideally we would want to utilize this information in the resulting hardware, for example extracting $(x+1)\times(x+1)$ as $\texttt{let }y=x+1\text{ in }y\times y$. However, this makes the extraction problem an inherently global problem over the e-classes, in the sense that the optimal e-node implementation for a given e-class may depend on the selected e-node implementation of the other e-classes in the graph. Previous solutions have solved this by posing optimal extraction as an integer linear programming (ILP) problem [13, 14], and we follow the same approach in this work.

Using the notation defined in Section II-B, for each e-node $n\in\mathcal{N}$ we associate a cost, $\text{cost}(n)$, given by the theoretical cost function, and a binary variable $x_{n}\in\{0,1\}$, which indicates whether $n$ is implemented in the final extracted RTL. Our objective is to minimize the total implementation cost, as described by (1). (2) then guarantees that for every node $n$, we implement a node from each of its child e-classes. Lastly we introduce $\mathcal{S}$, the set of e-classes representing the desired expressions to implement in RTL. (3) then ensures that all these outputs are produced by the final RTL.

E-graphs may contain cycles, e.g. $x+0\to x$ induces a cycle. By introducing a topological sorting variable, $t_{c}$ for each e-class $c$, and associated constraints (4) where $N$ is the number of e-classes and $\mathcal{C}(n)$ is the e-class containing node $n$, we ensure that the output expression is acyclic.

If we select a node $n\in\mathcal{N}_{c}$ with child $k$, i.e. $x_{n}=1$, this constraint simplifies to $t_{c}\geq t_{k}+1$ to get a topologically sorted result, whereas in the case $x_{n}=0$, the constraint is vacuously satisfied. We use the open source GLPK solver to calculate solutions to this ILP [21].

We deploy the ILP extraction method in cases where sharing common subexpressions offers some improvement, otherwise for improved performance we can resort to the standard egg extraction method [5].

The first benchmark is a kernel from a media module, which is an industrially relevant example supplied by Intel. It was manually optimized by a single engineer over the course of one week. Our tool automatically matches the results obtained via manual optimization, making use of the “Merge Mult Array” rewrite. Figure 4 shows the area-delay curves for the original and optimized architectures. The optimized design performs strictly better, reducing the minimum achievable delay by 13% with a 28% area reduction.

Next we have a set of benchmarks taken directly from [4], which are intended to show how this prior work is generalized by our optimization tool. The first example is a finite-impulse response (FIR) filter with 8 taps. The second example, a computational kernel of the ADPCM decoder [22], is an approximation to a 16$\times$4 multiplier. These two examples are optimized by deploying the “Arithmetic Logic Exchange” class of rewrites described in Section III-B, to cluster additions together. The FIR filter example is particularly interesting as it also introduces a multiple constant multiplication (MCM) problem [7]. Since the ILP formulation of extraction correctly accounts for the cost of common subexpressions, the minimal area solution it extracts maximally shares common subexpressions, to generate $\{2\times x,3\times x,...,7\times x\}$. For example, $2\times x$ and $3\times x$ are constructed as follows

Looking more generally at the MCM problem, using the “Constant Expansion” rewrites, we are able to match the operator count of a solution from [6]. An example from this paper generates adder graphs to compute $\{3\times x,7\times x,21\times x\}$. The optimal design from our tool maximally shares common subexpressions to generate the data-flow graph shown in Figure 5, which uses 3 addition/subtraction operations to generate the results. Since logic synthesis likely implements constant multiplication using a CSD representation [8], sharing common subexpressions generates a $15\%$ larger circuit than the basic architecture for small delay targets as the design does not match the performance of CSD.

Finally we consider a shifted FMA, which discovers the opportunity to merge a multiplication and addition that were originally separated by a left shift $(a\times b)<<S+c$. Using the “Left Shift Mult” and “FMA Merge” rewrites this can be implemented as an FMA block. The approach proposed by Verma, Brisk and Ienne misses this opportunity since it will not move the shift over the multiplication [4].

The last set of results considers parameterizable RTL, as in general, RTL engineers do not generate alternative architectures for different bitwidths. As a result, designers may be sacrificing performance. To demonstrate this we again consider the FIR filter, but the 4 tap variant, where the original architecture is shown in Figure 6, which we will refer to as Architecture 0. We optimized Architecture 0 using our tool for increasing bitwidths and observed that the selected architecture varied between Architectures 0, 1 and 2 (Figure 6). The difference between Architectures 1 and 2 is that the addition involving $Z4$ in Architecture 1 has been pushed back over the right shift, incurring the cost of an extra left shift, but saving a full carry-propagate adder.

The architectural choices are determined by the theoretical two-input gate cost metric, discussed in Section III-C. We synthesized each of the three architectures for bitwidths $4,8,...64$, updating the delay target each time to the minimum delay that all three architectures could meet at that bitwidth. The theoretically optimal architecture for each bitwidth can be seen in Figure 6. For 56% of the bitwidths, the architecture selected by the theoretical cost metric gave the lowest area result from logic synthesis. The ability to always choose the minimum area design according to logic synthesis using the theoretical metric is limited, since logic synthesis always incorporates delay considerations into its results and there is noise in the results. We demonstrate this ‘noise floor’ in Section V.

For these test cases only the “Shifted FMA” e-graph saturated as we limited the exploration to 10 iterations of e-graph rewriting [5]. For test cases using ILP extraction we limited the solver to a 100 second time budget, which means that the FIR Filter solution is classified as feasible but is possibly sub-optimal. There is a tradeoff between growing the e-graph and solving the ILP efficiently. Only the MCM solution included the “Constant Expansion” rewrite class since these rewrites lead to exponential growth of the e-graph.