Applying egg to SSA-based IRs

[femtomc]

Hi! What are the prospects of utilizing egg to perform rewrite-driven transformations on SSA-based IRs (?)

I read this paper (whose authors appear to be collaborators with y'all): https://www.cs.cornell.edu/~ross/publications/eqsat/ where they describe E-PEG-based program representation.

On the other hand, I'm working on IRs which are closer to strict SSA (where all control flow is de-sugared to basic blocks and branching) and everything is linear (e.g. blocks are vectors of instructions) vs. strictly tree-like.

I'd love to find or understand how equality saturation could be applied to this representation structure -- it's likely that maybe this information can be gleaned from the above paper (but I expect that maintainers here might be able to unpack this more than I can).

[chandrakananandi]

Hi! I use egg for a language which is mostly-SSA. In my case, things "just work" so far because I mostly care about straight line blocks. The only additional information I track which may be necessary for other SSA-based usages as well is to make sure you don't use variables before they are defined. This, in my case, was possible with a simple e-class analysis.

[femtomc]

I’m curious about this — are you working with single blocks at a time? How does scope work in your language — e.g. can blocks which dominate provide SSA references which do not need to be passed as block arguments (or otherwise represented in phi nodes)?

e.g. ignoring flow between blocks, I can sort of imagine just restricting attention to def-use chains, and then you have a nice DAG-like thing, and there’s probably some details about fixing the chains up after rewriting.

Re — is this project open source at all? I’m just getting started with egg so I’m curious about usage like this. No worries if not!

[chandrakananandi]

My use case at the moment involves working with a single block at a time which are all jump-free. I don't actually modify the code itself (it's part of a verification process, so I shouldn't :) ), but I generate a summary of linear invariants from the blocks. Your question about the dominating block is interesting --- I will need to run some more experiments to see what happens in that case. I will get back to you on this once I have more info.

Unfortunately the project is not open-source yet but I would like to make it available in some form (maybe an article or blog post) soon.

[mwillsey]

Hi, I think @chandrakananandi is right. I don't think there is any major technical blocker stopping you from doing this today!

The original eqsat paper you linked to contains a lot of good ideas on how to encode your problem into an e-graph. If I'm understanding your setting, you'll still want E-PEGs if you want to encode loops in a transparent way that you can optimize through. If you only care about optimizing one basic block (one DFG) at a time, then you don't need them.

The egg paper (and tool) are mostly innovating in how equality saturation is done. How you encode your problem is still up to you!

Despite there not being any huge blockers, I still think it's a large and challenging task, and one that I'd like to try to tackle at some point (if I can find the time), or see someone else take a stab at!

One thing that I will add: you'll have a much easier time with a tree- or dag-like IR that a linear, mutating IR. All the rewrites that you do are over trees or dags, and just overall the e-graph doesn't do a lot for you if your language is heavily sequential. If that's your case, consider building def-use chains or some other method to make things more graph-like for the e-graph.

[femtomc]

That’s a great idea! Thanks Chandra!

[Bastacyclop]

Regarding types, an alternative to using e-class analyses is to embed the types in the e-graph (the easiest with egg would probably be to include types in terms/e-nodes, in my own project I actually include types in e-classes). Then, it is possible to write untyped (or partially annotated) rewrite rules and to implement rewrite rule type inference: producing a typed rule from the untyped one. A typed rule can match (deconstruct) types on its left-hand-side, and construct types on its right-hand-side. It is also possible for rewrite rule type inference to generate appliers on the right-hand-side if purely syntactic constructs are not sufficient. I am successfully using this approach for a typed lambda calculus language.

[michel-steuwer]

You can read about @Bastacyclop's work in our paper that you find on arXiv: https://arxiv.org/pdf/2111.13040.pdf

We are very interested in bringing equality saturation to MLIR and we have a couple of contacts at Google working with MLIR who are very interested in this as well.

[femtomc]

@michel-steuwer Cool! My project is not really affiliated with MLIR in any formal way -- I'm just inspired by their IR design, and I'm exploring some PL questions associated with the design concepts. The way I'm experimenting with things now, I actually split types out of the IR structure -- and I'm putting them into an abstract interpretation framework. I think that means that it will be harder in my case to get the right information to the right places (at least, it requires a bit more thought).

[femtomc]

(on the other hand, the reason why I reached out to egg is because my experiment is written in Rust, so it seemed easier to integrate now, versus whatever is required to provide the same capabilities to C++)