4 ms·

Automated Code Optimization with E-Graphs

102 points by matt_d 5y ago

gugagore 5y ago
Here is a short (5 min) video explaining E-Graphs: https://www.youtube.com/watch?v=ap29SzDAzP https://www.youtube.com/watch?v=ap29SzDAzP There is also a longer (30 min) video that goes into details. I recommend watching them! Also consider this blogpost: https://blog.sigplan.org/2021/04/06/equality-saturation-with-egg/ https://blog.sigplan.org/2021/04/06/equality-saturation-with...
- minetest2048 5y ago
  > https://www.youtube.com/watch?v=ap29SzDAzP https://www.youtube.com/watch?v=ap29SzDAzP This video isn't available anymore
  - gugagore 5y ago
    Oops, sorry: https://www.youtube.com/watch?v=ap29SzDAzP0 https://www.youtube.com/watch?v=ap29SzDAzP0
ChrisRackauckas 5y ago
Awesome to see this here! A lot of this work was done by an undergrad as part of his thesis, and I summarized the work in this Twitter thread along with our planned directions for 2022 (https://twitter.com/ChrisRackauckas/status/1477274812460449793 https://twitter.com/ChrisRackauckas/status/14772748124604497...). Though the abstract also captures the idea quite well too: Abstract: This thesis proposes an advanced, generic and high-level code rewriting and analysis system in the Julia programming language, providing applied equality saturation in the presence of multiple dispatch and metaprogramming. We show how our system can practically solve some challenging problems: Can programmers implement their own high-level compiler optimizations for their domain-specific scientific programs, without the requirement of them being compiler experts at all? Can these optimizers be implemented by users in the same language and inside the same programs they want to optimize, solving the two-language problem? Can these compiler optimizers be written in a high-level fashion, as equations, without the need to worry about the rewriting ordering? Thus, can symbolic mathematics do high-level compiler optimizations or vice-versa?
DannyBee 5y ago
For those who are not compiler experts, let me try to give a simplified (but not really wrong) explanation. About 10 years ago, a paper came out on using something called equality saturation to do compiler optimization. This is a fancy term that really means "generate all equivalences in the IR rather than destructive IR updates". IE if you have a = 5 + 5 "Destructive" optimizers will rewrite this to: a = 10 "Equality saturation" would simply annotate this with <a = 10> a = 5 + 5 Rather than destroying the original program. Each optimization generates sets of equivalences and equalities, rather than updating the IR. Over time, equality saturation has moved to more formally expressed rewrite systems as the optimizers. Something eventually comes through and picks the most "optimal" set of equivalent code for the program. Note that almost all compilers already do this in various phases for various reasons - LLVM and GCC both, for example, annotate the IR with inequalities and other equivalences prior to doing value optimizations, etc. Equality saturation has it's niceties: 1. phase ordering is less of an issue- assuming you do this well, you will have no missed optimizations due to alternative representations that were not chosen. 2. more formalized rewrite systems are often verifiable, simpler, and have less bugs It also has downsides: 1. lots of memory 2. very complicated exploration of the optimization space - none of the optimizations know which set of equivalences will be chosen in the end, so if they want the best code, they have to explore optimizing all the equivalences that have been generated before. 3. Most rewrite passes (formally expressed or not) individually are not optimal enough (and some can be proven impossible to make optimal) that you will still end up with phase ordering issues, just less of them, whether you express your passes as formal rewrite rules or not. In general, most of these problems end up as NP-hard constraint solving problems or equivalent. Truthfully, it's an approach that is interesting (I went down this rabbit hole myself in LLVM at one point), and if you applied good ML or heuristics to it, i bet you could find a reasonable tradeoff of what to explore vs not, and get better results than the compiler gets now. Expressing them as formally verifiable rewrite rules rather than ad-hoc passes will also make for a lot less compiler bugs ;) Note that the general notion of expressing optimizations in a equivalence/rewrite language and applying those iteratively, is not new (Stratego was an example of such a thing, there are many). Lots of them were destructive, there were some non-destructive ones, but none (IMHO) as well formalized as equality saturation or e-graphs. This is in part because the computational cost has historically been quite high. In any case, fast forward 12 years, this paper generalizes equality saturation to a turing complete graph expression (the original paper used extended PEGs) and language for doing rewrites, and uses it to optimize Julia programs, and provide a way for folks to write their own rewrite/equality/etc rules. This is awesome work. It still has the same issues (matching is NP-hard, etc). That is, as i said, likely solvable in a practical way by better guiding matching/rewriting/equivalence finding. You can prove you can't get optimal answers anyway (it would require the programs to be statically decidable, and none are), so no point in trying. The basic problem with both approaches, at least historically is that it's not better enough for the cost. Those who know me know I am a total sucker for abstractly nice optimization algorithms that produce really great results. I'm a big believer we've thrown away a lot of years of compiler theory in the pursuit of making "good enough" compilers. Lots of knowledge and theory is ignored when it would make things a lot better. But even I have my limits, and pragmatically, like I said, these kinds of approaches have just never been better enough for the cost, nor have they caught on among compiler folks yet for their ability to have less buggy compilers (that part is sad, FWIW) This is not a knock on the paper - it's a bachelors thesis, and a good one at that. But in terms of practicality of these approaches in real compilers, i'll believe it when i see it :)

102 points by matt_d 5y ago

Here is a short (5 min) video explaining E-Graphs: https://www.youtube.com/watch?v=ap29SzDAzP https://www.youtube.com/watch?v=ap29SzDAzP There is also a longer (30 min) video that goes into details. I recommend watching them! Also consider this blogpost: https://blog.sigplan.org/2021/04/06/equality-saturation-with-egg/ https://blog.sigplan.org/2021/04/06/equality-saturation-with...

minetest2048 5y ago
> https://www.youtube.com/watch?v=ap29SzDAzP https://www.youtube.com/watch?v=ap29SzDAzP This video isn't available anymore
- gugagore 5y ago
  Oops, sorry: https://www.youtube.com/watch?v=ap29SzDAzP0 https://www.youtube.com/watch?v=ap29SzDAzP0

ChrisRackauckas 5y ago

Awesome to see this here! A lot of this work was done by an undergrad as part of his thesis, and I summarized the work in this Twitter thread along with our planned directions for 2022 (https://twitter.com/ChrisRackauckas/status/1477274812460449793 https://twitter.com/ChrisRackauckas/status/14772748124604497...). Though the abstract also captures the idea quite well too: Abstract: This thesis proposes an advanced, generic and high-level code rewriting and analysis system in the Julia programming language, providing applied equality saturation in the presence of multiple dispatch and metaprogramming. We show how our system can practically solve some challenging problems: Can programmers implement their own high-level compiler optimizations for their domain-specific scientific programs, without the requirement of them being compiler experts at all? Can these optimizers be implemented by users in the same language and inside the same programs they want to optimize, solving the two-language problem? Can these compiler optimizers be written in a high-level fashion, as equations, without the need to worry about the rewriting ordering? Thus, can symbolic mathematics do high-level compiler optimizations or vice-versa?

DannyBee 5y ago

For those who are not compiler experts, let me try to give a simplified (but not really wrong) explanation. About 10 years ago, a paper came out on using something called equality saturation to do compiler optimization. This is a fancy term that really means "generate all equivalences in the IR rather than destructive IR updates". IE if you have a = 5 + 5 "Destructive" optimizers will rewrite this to: a = 10 "Equality saturation" would simply annotate this with <a = 10> a = 5 + 5 Rather than destroying the original program. Each optimization generates sets of equivalences and equalities, rather than updating the IR. Over time, equality saturation has moved to more formally expressed rewrite systems as the optimizers. Something eventually comes through and picks the most "optimal" set of equivalent code for the program. Note that almost all compilers already do this in various phases for various reasons - LLVM and GCC both, for example, annotate the IR with inequalities and other equivalences prior to doing value optimizations, etc. Equality saturation has it's niceties: 1. phase ordering is less of an issue- assuming you do this well, you will have no missed optimizations due to alternative representations that were not chosen. 2. more formalized rewrite systems are often verifiable, simpler, and have less bugs It also has downsides: 1. lots of memory 2. very complicated exploration of the optimization space - none of the optimizations know which set of equivalences will be chosen in the end, so if they want the best code, they have to explore optimizing all the equivalences that have been generated before. 3. Most rewrite passes (formally expressed or not) individually are not optimal enough (and some can be proven impossible to make optimal) that you will still end up with phase ordering issues, just less of them, whether you express your passes as formal rewrite rules or not. In general, most of these problems end up as NP-hard constraint solving problems or equivalent. Truthfully, it's an approach that is interesting (I went down this rabbit hole myself in LLVM at one point), and if you applied good ML or heuristics to it, i bet you could find a reasonable tradeoff of what to explore vs not, and get better results than the compiler gets now. Expressing them as formally verifiable rewrite rules rather than ad-hoc passes will also make for a lot less compiler bugs ;) Note that the general notion of expressing optimizations in a equivalence/rewrite language and applying those iteratively, is not new (Stratego was an example of such a thing, there are many). Lots of them were destructive, there were some non-destructive ones, but none (IMHO) as well formalized as equality saturation or e-graphs. This is in part because the computational cost has historically been quite high. In any case, fast forward 12 years, this paper generalizes equality saturation to a turing complete graph expression (the original paper used extended PEGs) and language for doing rewrites, and uses it to optimize Julia programs, and provide a way for folks to write their own rewrite/equality/etc rules. This is awesome work. It still has the same issues (matching is NP-hard, etc). That is, as i said, likely solvable in a practical way by better guiding matching/rewriting/equivalence finding. You can prove you can't get optimal answers anyway (it would require the programs to be statically decidable, and none are), so no point in trying. The basic problem with both approaches, at least historically is that it's not better enough for the cost. Those who know me know I am a total sucker for abstractly nice optimization algorithms that produce really great results. I'm a big believer we've thrown away a lot of years of compiler theory in the pursuit of making "good enough" compilers. Lots of knowledge and theory is ignored when it would make things a lot better. But even I have my limits, and pragmatically, like I said, these kinds of approaches have just never been better enough for the cost, nor have they caught on among compiler folks yet for their ability to have less buggy compilers (that part is sad, FWIW) This is not a knock on the paper - it's a bachelors thesis, and a good one at that. But in terms of practicality of these approaches in real compilers, i'll believe it when i see it :)