Representing Loops within Egg

[Kixiron]

I'm really struggling to figure out how to express more advanced control flow with egg, things like for and while loops. From exploring the problem space I feel like more traditional compiler constructs like lowering to tail-controlled do ... while ... loops and some sort of phi node are the most promising candidates, but I'm still not entirely sure how to do this effectively.
Modeling loops in a form that egg understands is something critical to whole-program optimization, autovectorization and other fancy types of optimization and I'm eager to get it working, thanks!

[Bastacyclop]

The original paper on equality saturation (https://www.cs.cornell.edu/~ross/publications/eqsat/) features loops, I would assume that you can re-use the same loop representation technique with egg?

[mwillsey]

If you want to create cycles, you can first insert an acyclic term and then create a cycle with a union. So if you insert f(a) and then union that with a, you'll get a cycle. In your case, you'll just end up adding a "dummy" node.

Extraction will never extract a term with a cycle in it, so I think that's the issue.

[Kixiron]

Yep, that's what the backtrack refers to (I should have clarified that more): I added a tombstone enode into the egraph as a placeholder, unioned it with the final node and then removed the tombstone from the egraph, then this panic fired off when trying to extract the cyclic egraph (interestingly it always chooses the tombstone if given a choice between it and the cycle, even if it's given an absurdly high cost, I assume that's by design though). Is there any way to actually extract cyclic terms from the egraph or plans to implement that (or pointers on implementing it myself)?

Here's the full egraph:

0: [UInt(10)]
1: [UInt(1)]
2: [Other("phi", [0, 3])]
3: [Add([2, 13]), Add([13, 2]), Sub([2, 1])]
5: [UInt(0)]
6: [Eq([2, 5])]
7: [Break]
8: [OpaqueFunc(Func { id: 0 })]
9: [Unit]
10: [Apply([8, 9])]
11: [If([6, 7, 10])]
12: [Loop(11)]
13: [Mul([1, 16]), Mul([16, 1]), Neg(1)]
16: [Int(-1)]

[mwillsey]

It will always choose the tombstone because a cyclic term has "infinite" cost from the perspective of our bottom-up extractor. There is no way to coerce our extractor to give you a cyclic term. The original PEG paper uses ILP to extract cyclic terms from their e-graphs, and the same could work here. I have no doubt that a simpler approach might work as well, but I haven't done. If you come up with something (and you'd need a way to represent cyclic terms as well), I'd love to hear about it!

[Kixiron]

So extracting cyclic terms would require writing my own extraction infrastructure? And do you have a link to the PEG paper you're referencing?

[mwillsey]

Yes, you'll need to write your own. Here is the paper: https://dl.acm.org/doi/pdf/10.1145/1480881.1480915

[jameysharp]

I'm grateful that you mentioned RVSDG here, @Kixiron. Based on your suggestion I read those papers this month, especially "Perfect Reconstructability of Control Flow from Demand Dependence Graphs", which convinced me that it's a good model for arbitrary control flow in an e-graph. Re-reading my own blog post from five years ago ("Search-based compiler code generation"), I'd apparently discovered the "Perfect Reconstructability" paper then, but by the time I learned about equality saturation I'd forgotten all about it.

And now I have a working prototype of e-graphs for RVSDG, based on egg. Every valid RVSDG is acyclic, so avoids the extraction difficulties discussed here, but you can still round-trip even irreducible control flow through it. I'm very pleased both with egg and with how this experiment has turned out. And yes, to the original question, loops seem to work fine in this model.

https://github.com/jameysharp/optir/

[mwillsey]

Hello Jamey!

This is very cool work! I shared this with the other folks in the egg research group, and we all agree that this is a great response to the very common question of "what about loops/control flow?". Keep us posted if you end up using this in a larger setting!