Skip to content

Qi Compiler Sync Nov 4 2022

Siddhartha Kasivajhula edited this page Dec 15, 2022 · 6 revisions

E-Graph Special

Qi Compiler Sync Nov 4 2022

Adjacent meetings: Previous | Up | Next

Summary

We discussed e-graphs! Oliver Flatt from the Herbie/Egg team was kind enough to join us and share some thoughts and advice.

Due to a Daylight Saving scheduling debacle, the event scheduled at 10am Pacific time appeared as 11am on Discord (the event was created in Mexico which has already "fallen back" while the US hasn't fallen back yet... so the time difference now is less than it was at event creation time or something. I plead "nobody understands time zones" and "fewer still understand daylight saving time"), so we had two cohorts of people join, and the second served as a kind of breakout session ("You're not going to believe this, but the meeting just ended 😅" - Sid to second cohort). Yet another attendee forgot about the meeting and is going to catch up by reading these meeting notes, and will therefore be considered the third cohort.

Background

Last time, Micheal brought up e-graphs as a potentially interesting data structure to consider for structuring compiler passes in an optimal way. Ben brought this up on Discord and then one thing led to another and before we knew it, it had turned into a giant e-graphs party. Evidently, a lot of people are interested in this data structure and its potential to help with optimizing compilers and related technologies, so this time's meetup was geared towards understanding e-graphs and what provisions Racket has for working with them.

E-graphs

An E-graph is a data structure that represents equivalences between terms of a language, allowing selection of "optimal" terms with respect to an arbitrary cost function. The storage of terms is efficient in that there are no duplicates. Additionally, the computation of equivalence on adding fresh term rewrite rules is done in such a way that it preserves "congruence," that is, it infers equivalence of terms in the entire data structure that are implied by such newly declared equivalences.

More detailed and clear descriptions can be found in the Egg tutorial.

E-graphs could be useful for writing optimizing compilers, among other things.

Egg, Regraph, and Egglog

Oliver summarized the options and tradeoffs in the Racket ecosystem this way:

  1. egg with a text/ffi interface -- Well supported, more features, inconvenient FFI
  2. Regraph -- Racket, easier to implement, fewer features
  3. egglog -- Better than egg, text interface, super alpha

The FFI is inconvenient because (in general) it requires attention to memory management and other such details that Racketeers would prefer to avoid. On the other hand, Egg includes a number of features that aren't present in Regraph, and would take (presumably significant) effort to implement, including:

  • optimizations for program analysis
  • optimizations for the extraction step (using Iterative Linear Programming (ILP) - a way to efficiently solve special cases of NP-hard problems)
  • proofs for equivalences

Egglog combines Egg and Datalog to achieve better results. For example, one of the cases that Egg struggles with is division by zero. It would be able to make stronger inferences in many cases if it could assume that division by zero is not a possibility. In order to improve this, Egglog attempts to use Datalog as a preprocessing step to segregate the input into cases where the division-by-zero possibility is present vs eliminated, and then Egg is able to achieve better results by starting from this stronger set of base assumptions.

Practical Considerations for Projects Using E-graphs

E-graphs need not be used wholesale for all of the features they provide. They could be used purely as a data structure for storing equivalences while using a different mechanism for measuring performance and for "extraction," (i.e. selecting an optimal term with respect to a cost function) or they could be used for storing equivalences for specialized subsets of the program space -- i.e. small e-graphs for special cases instead of a large e-graph for everything.

"No one's bitten the bullet" on using e-graphs for a real-world compiler, due to the increase in compile time. So if compile time is a significant metric, e-graphs may not be the best approach (yet).

The one exception is the Cranelift project which made a number of alterations in order to use e-graphs (e.g. "acyclic e-graph").

Other cases where e-graphs and Egg have been used include floating point math and EVM blockchain programs. It's unclear whether they have been used much in functional languages (like Qi).

In practice, "context is a problem" and is one of the first things to figure out. [TODO: don't recall what this refers to - context in what sense? Maybe this is related to "datalog modulo theories" mentioned below]

A Rackety Interface to E-graphs?

From a Racket perspective, an ideal interface (suggested by Jack Firth on Discord) would be to:

  • provide syntax objects as the nodes of the e-graph ("e-nodes"), and
  • specify rewrite rules as macros

It may be possible to have such an interface, especially if we don't need the program analysis component of what Egg does.

Program analysis seems to be a way to determine which rewrite rules are applicable in a way that isn't purely based on matching a pattern. These analyses are composable with the rewrite rules themselves. Since they aren't pattern-matching based, it may not be a straightforward mapping to Racket macros.

The Promise of E-graphs

In theory, e-graphs could allow us to avoid a fixed ordering of compiler passes and instead take any order that would result in the optimal outcome, as determined during the actual compilation process as opposed to being hardcoded beforehand. The big question is how to define the cost function so that we are able to recognize these optimal expressions. Naively with compiler passes and rewrite rules, the rules encode our assumption that the right-hand side of the rule is more performant than the left-hand side, and the ordering of compiler passes allows us to prioritize more-optimal rewrite rules ahead of less-optimal ones. But there is always the chance that we get this wrong, and more importantly, that it may be impossible to structure the passes to always arrive at the most optimal result, as the "rewrite multiplication-by-two to left-bitshift-by-one" example in the Rust documentation shows. That is, this particular optimization is indeed faster, but it could land us on a local optimum and prevent further simplifications and optimizations that would have been possible by application of other rules (such as recognizing that x * 2 / 2 = x). E-graphs allow us to apply all rules in every order so that the global optimum can always be found.

... which brings us back to the cost function. Is it truly possible to define a cost function that can efficiently recognize optimal versions? "Shorter generated Racket code," for example, is a crude heuristic when we consider that loop unrolling generates more efficient code by generating more code. Given this difficulty of assessing optimality, could the cost function perform better in practice than manually encoding these assumptions as compiler passes and possibly ending up in local maxima? And even so, would it be performant enough (i.e. at compile time) to be worth it?

Other Things Being Explored in the E-graph space

  • "Datalog modulo theories" - employing datalog in well-bounded domains like "math" and "Lisp" rather than in a context-free way.
  • Canonical representations -- e-graph is a purely syntactic representation. It might be useful to explore more canonical ("prolog-like"? May be related to "constrained Horn clauses"?) representations.

Next Steps

  • It could be nice to evaluate using e-graphs for a small set of rewrite rules (either via regraph or an FFI)
  • If you use e-graphs in your next project, please share your experience!

Resources

Attendees

Cohort 1: Chris, Hans, Jair, Jason, Michael, Oliver, Sid

Cohort 2: Julia, Ryan, Sid

Cohort 3: Ben 🌄 🛣 🚗

Clone this wiki locally