Use rewrites to convert relay to glenside

[gussmith23]

This is a long-needed change which implements the Relay-to-Glenside compiler using egg rewrites.

In the past, the Relay-to-Glenside compiler would read in Relay programs using the Relay Rust bindings, and then construct Glenside expressions out of those expressions. There were no Relay expressions in Glenside at the time, and so everything needed to be converted directly to Glenside. It was only later that we began adding Relay expressions into Glenside. Now, all Relay constructs of interest are explicitly representable in Glenside using RelayOperatorCall and other constructs. With that in mind, the cleaner way to implement a Relay to Glenside compiler would be by importing a Relay program as RelayOperatorCalls, and then using rewrites to rewrite the various subexpressions to their Glenside equivalents. We can implement the Glenside-to-Relay compiler the same way. This has multiple benefits:

1. If we implement the Glenside-to-Relay compiler the same way, then we can run Relay-to-Glenside, Glenside-to-Glenside, and Glenside-to-Relay rewrites all at the same time, simplifying the entire pipeline significantly.
2. We don't have to choose a single Glenside representation for any Relay operator. Previously, the Relay-to-Glenside compiler needed to choose a way to represent, for example, conv2d, which can be expressed in a few ways in Glenside. Now, to cover the multiple ways to represent a complex Relay operator in Glenside, we need only include multiple rewrites.
3. Equivalent Glenside and Relay nodes are automatically unified. Currently, Mike has some code which manually unifies each Relay node with its equivalent Glenside node, after Relay-to-Glenside compilation. This is totally unnecessary if we use rewrites, as rewrites do this exact unification. Once Mike had to implement this code, it was pretty immediately obvious that we needed to rethink our compilation method.

My plan for doing this:
for each Relay operator supported in Glenside, add a rewrite which rewrites from a Relay operator invocation to a Glenside expression. How should we test, however?

Operators remaining to be implemented:

• conv2d
• conv1d
• softmax
• relu
• sqrt
• negative
• maxpool2d
• global_avg_pool2d
• expand_dims
• pad
• dense
• add
• multiply
• divide
• maximum (we only implement this op opaquely in Glenside at the moment)
• minimum (we only implement this op opaquely in Glenside at the moment)
• batch_flatten
• bias_add
• concatenate
• reshape
• transpose
• squeeze

[gussmith23]

Current problem I'm facing: when we go to add, for example, the conv2d expression we construct with the conv2d helper function from from_relay, we can't get back the Id of the glenside conv2d expression; instead we get back the id of the conv2d RelayOperatorCall. This is purely because the add function adds everything in a RecExpr to the egraph and returns the very last thing added, which in this case will be the RelayOperatorCall. An easy hack to get around this w/o changing too much is to make sure the glenside conv2d expression is the very last thing added in the from_relay::conv2d() helper function.

In the future, this won't be an issue, as the helper function will no longer add RelayOperatorCalls.

[gussmith23]

Running into a bug while implementing the rewrite for conv2d. I think it has to do with how I'm using the existing from_relay::conv2d() helper function to construct a RecExpr, and then adding that RecExpr to the egraph. Specifically, the helper function requires you to specify the Ids of the existing data/weight expressions in the RecExpr. However, when I call the helper function, I pass in Ids from the egraph. These are not valid Ids in the RecExpr. Then, when the egraph tries to add the RecExpr, it uses that invalid Id to do an invalid lookup, and things go off the rails from there. I think the solution might be to try and use the helper function to construct a pattern (which is just a special RecExpr which can have Vars at the leaves) and use the new add_instantiation function to add the pattern to the egraph.

[gussmith23]

I need to update egg to use add_instantiation. I'm seeing if I can just update egg to vanilla 0.7.1 without Mike's changes and get away with it. Otherwise, we can add the add_instantiation impl on top of Mike's changes, or (harder way) rebase Mike's changes on top of 0.7.1.

[gussmith23]

Okay, conv2d rewrite is tentatively working. Who knows what else I broke in the process, though.

[gussmith23]

I added add_instantiation manually. We should update egg when we get the chance. Specifically, we should stop using a fork of egg; I'm not sure it's needed!

[gussmith23]

I'd like to use the existing test machinery to test the new code. Specifically, the test! macro; it would be nice to convert it to be usable in the new setting. Then we could pull over our tests with ease. To do this (and just in general), we need to convert the from_relay function to parse Relay programs and insert Relay constructs into the egraph (rather than Glenside constructs). Then we would have the extra step of running the egraph with the correct Relay-to-Glenside rewrite and checking for the expected Glenside program.

[gussmith23]

I'm working on this now. I'm making it so that the from_relay function adds just a RelayOperatorCall to the expression, and nothing more. Then, I'll make a new version of the test macro which is nearly the same, except it accepts a list of rewrites to run over the egraph (which do the compilation).

I'm doing this for conv2d first.

[gussmith23]

Seems to be working for conv2d. Great! Now we can proceed by:

For a given Relay operation (e.g. conv1d):

1. Create the rewrite for the operation by making a rewrite which rewrites from a RelayOperatorCall to the Glenside implementation of the operator.
2. Change from_relay() so that it only constructs a RelayOperatorCall, and does not construct an actual Glenside expression.
3. Copy the test! (or test!s -- there may be multiple) for that operation from from_relay/mod.rs to rewrites.rs.
4. Add the necessary rewrites to the test!s.

The tests should work out of the box!

[gussmith23]

conv1d done, proceeding as normal. had to add some fields to conv1d. We make a lot of assumptions about what kinds of conv1ds we can handle, and those are still explicitly set out in from_relay(), but they're not necessarily made explicit anywhere else.

[gussmith23]

I'm having an issue writing a test for expand_dims, as I can't write a string for a pattern that will match an Int64 literal. it just parses as a usize. I'm not sure what the right way to fix this is. perhaps just convert everything to int64? But expand_dims is working, fwiw.

[gussmith23]

I'm having an issue with the concatenate rewrite. Specifically, I need an additional rewrite which simplifies (tuple-get-item (construct-tuple ?a ...) 0) --> ?a. I have written the rewrite, and it tentatively works. But currently, when I run the rewrite (with the concatenate relay->glenside rewrite), the concatenate rewrite fires, introduces a (tuple-get-item (construct-tuple ...) ...) pattern, but then the tuple-simplification rewrite doesn't fire for some reason. Specifically, the rewrite runner claims that things saturate, which doesn't make sense to me.

[gussmith23]

Okay, the issue seemed to be that I was doing a union() inside of the applier for the concatenate rewrite, and that was breaking things. The problem is that the egg API updated and soon enough i think we will need to union() in the applier. So i'll have to make sure I update the code when I update egg.

[gussmith23]

All operators implemented, unless I missed some!

[gussmith23]

I'm going to go ahead and merge this in, as I just want to move on with it!