Crate cranelift_egraph

Expand description

ægraph (aegraph, or acyclic e-graph) implementation.

An aegraph is a form of e-graph. We will first describe the e-graph, then the aegraph as a slightly less powerful but highly optimized variant of it.

The main goal of this library is to be explicitly memory-efficient and light on allocations. We need to be as fast and as small as possible in order to minimize impact on compile time in a production compiler.

The e-graph

An e-graph, or equivalence graph, is a kind of node-based intermediate representation (IR) data structure that consists of eclasses and enodes. An eclass contains one or more enodes; semantically an eclass is like a value, and an enode is one way to compute that value. If several enodes are in one eclass, the data structure is asserting that any of these enodes, if evaluated, would produce the value.

An e-graph also contains a deduplicating hash-map of nodes, so if the user creates the same e-node more than once, they get the same e-class ID.

In the usual use-case, an e-graph is used to build a sea-of-nodes IR for a function body or other expression-based code, and then rewrite rules are applied to the e-graph. Each rewrite potentially introduces a new e-node that is equivalent to an existing e-node, and then unions the two e-nodes’ classes together.

In the trivial case this results in an e-class containing a series of e-nodes that are newly added – all known forms of an expression – but Note how if a rewrite rule rewrites into an existing e-node (discovered via deduplication), rewriting can result in unioning of two e-classes that have existed for some time.

An e-graph’s enodes refer to classes for their arguments, rather than other nodes directly. This is key to the ability of an e-graph to canonicalize: when two e-classes that are already used as arguments by other e-nodes are unioned, all e-nodes that refer to those e-classes are themselves re-canonicalized. This can result in “cascading” unioning of eclasses, in a process that discovers the transitive implications of all individual equalities. This process is known as “equality saturation”.

The acyclic e-graph (aegraph)

An e-graph is powerful, but it can also be expensive to build and saturate: there are often many different forms an expression can take (because many different rewrites are possible), and cascading canonicalization requires heavyweight data structure bookkeeping that is expensive to maintain.

This crate introduces the aegraph: an acyclic e-graph. This data structure stores an e-class as an immutable persistent data structure. An id can refer to some level of an eclass: a snapshot of the nodes in the eclass at one point in time. The nodes referred to by this id never change, though the eclass may grow later.

A union is also an operation that creates a new eclass id: the original eclass IDs refer to the original eclass contents, while the id resulting from the union() operation refers to an eclass that has all nodes.

In order to allow for adequate canonicalization, an enode normally stores the latest eclass id for each argument, but computes hashes and equality using a canonical eclass id. We define such a canonical id with a union-find data structure, just as for a traditional e-graph. It is normally the lowest id referring to part of the eclass.

The persistent/immutable nature of this data structure yields one extremely important property: it is acyclic! This simplifies operation greatly:

When “elaborating” out of the e-graph back to linearized code, so that we can generate machine code, we do not need to break cycles. A given enode cannot indirectly refer back to itself.
When applying rewrite rules, the nodes visible from a given id for an eclass never change. This means that we only need to apply rewrite rules at that node id once.

Data Structure and Example

Each eclass id refers to a table entry (“eclass node”, which is different than an “enode”) that can be one of:

A single enode;
An enode and an earlier eclass id it is appended to (a “child” eclass node);
A “union node” with two earlier eclass ids.

Building the aegraph consists solely of adding new entries to the end of this table of eclass nodes. An enode referenced from any given eclass node can only refer to earlier eclass ids.

For example, consider the following eclass table:


   eclass/enode table

    eclass1    iconst(1)
    eclass2    blockparam(block0, 0)
    eclass3    iadd(eclass1, eclass2)

This represents the expression iadd(blockparam(block0, 0), iconst(1)) (as the sole enode for eclass3).

Now, say that as we further build the function body, we add another enode iadd(eclass3, iconst(1)). The iconst(1) will be deduplicated to eclass1, and the toplevel iadd will become its own new eclass (eclass4).

    eclass4    iadd(eclass3, eclass1)

Now we apply our body of rewrite rules, and these results can combine x + 1 + 1 into x + 2; so we get:

    eclass5    iconst(2)
    eclass6    union(iadd(eclass2, eclass5), eclass4)

Note that we added the nodes for the new expression, and then we union’d it with the earlier eclass4. Logically this represents a single eclass that contains two nodes – the x + 1 + 1 and x + 2 representations – and the latest id for the eclass, eclass6, can reach all nodes in the eclass (here the node stored in eclass6 and the earlier one in elcass4).

aegraph vs. egraph

Where does an aegraph fall short of an e-graph – or in other words, why maintain the data structures to allow for full (re)canonicalization at all, with e.g. parent pointers to recursively update parents?

This question deserves further study, but right now, it appears that the difference is limited to a case like the following:

expression E1 is interned into the aegraph.
expression E2 is interned into the aegraph. It uses E1 as an argument to one or more operators, and so refers to the (currently) latest id for E1.
expression E3 is interned into the aegraph. A rewrite rule fires that unions E3 with E1.

In an e-graph, the last action would trigger a re-canonicalization of all “parents” (users) of E1; so E2 would be re-canonicalized using an id that represents the union of E1 and E3. At code-generation time, E2 could choose to use a value computed by either E1’s or E3’s operator. In an aegraph, this is not the case: E2’s e-class and e-nodes are immutable once created, so E2 refers only to E1’s representation of the value (a “slice” of the whole e-class).

While at first this sounds quite limiting, there actually appears to be a nice mutually-beneficial interaction with the immediate application of rewrite rules: by applying all rewrites we know about right when E1 is interned, E2 can refer to the best version when it is created. The above scenario only leads to a missed optimization if:

a rewrite rule exists from E3 to E1, but not E1 to E3; and
E3 is cheaper than E1.

Or in other words, this only matters if there is a rewrite rule that rewrites into a more expensive direction. This is unlikely for the sorts of rewrite rules we plan to write; it may matter more if many possible equalities are expressed, such as associativity, commutativity, etc.

Note that the above represents the best of our understanding, but there may be cases we have missed; a more complete examination of this question would involve building a full equality saturation loop on top of the (a)egraph in this crate, and testing with many benchmarks to see if it makes any difference.

Rewrite Rules (FLAX: Fast Localized Aegraph eXpansion)

The most common use of an e-graph or aegraph is to serve as the IR for a compiler. In this use-case, we usually wish to transform the program using a body of rewrite rules that represent valid transformations (equivalent and hopefully simpler ways of computing results). An aegraph supports applying rules in a fairly straightforward way: whenever a new eclass entry is added to the table, we invoke a toplevel “apply all rewrite rules” entry point. This entry point creates new nodes as needed, and when done, unions the rewritten nodes with the original. We thus immediately expand a new value into all of its representations.

This immediate expansion stands in contrast to a traditional “equality saturation” e-egraph system, in which it is usually best to apply rules in batches and then fix up the canonicalization. This approach was introduced in the egg e-graph engine ¹. We call our system FLAX (because flax is an alternative to egg): Fast Localized Aegraph eXpansion.

The reason that this is possible in an aegraph but not (efficiently, at least) in a traditional e-graph is that the data structure nodes are immutable once created: an eclass id will always refer to a fixed set of enodes. There is no recanonicalizing of eclass arguments as they union; but also this is not usually necessary, because args will have already been processed and eagerly rewritten as well. In other words, eager rewriting and the immutable data structure mutually allow each other to be practical; both work together.

M Willsey, C Nandi, Y R Wang, O Flatt, Z Tatlock, P Panchekha. “egg: Fast and Flexible Equality Saturation.” In POPL 2021. https://dl.acm.org/doi/10.1145/3434304 ↩

Macros

trace

Conditionally-compiled trace-log macro. (Borrowed from cranelift-codegen; it’s not worth factoring out a common subcrate for this.)

Structs

BumpArena

BumpSlice

A slice in an arena: like a BumpVec, but has a fixed size that cannot grow. The size of this struct is one 32-bit word smaller than BumpVec. It is copyable/cloneable because it will never be freed.

BumpVec

A vector of T stored within a BumpArena.

CtxHashMap

A HashMap that takes external context for all operations.

EClass

An EClass entry. Contains either a single new enode and a child eclass (i.e., adds one new enode), or unions two child eclasses together.

EGraph

An egraph.

Id

An eclass ID.

NodeIter

An iterator over all nodes in an eclass.