Bring some of 3la-pldi-push-main to main

[gussmith23]

This PR brings changes from our PLDI 2022 push into main. It does so by starting from our 3la-pldi-push-main branch and reverting/undoing/modifying a bunch of the changes. This allows me to preserve the commit history (giving proper contribution attribution to Vishal and Mike!) while not merging everything from the PLDI push.

For each of the following files, we need to inspect the diff between 3la-pldi-push-main and main, and revert any changes that we don't want. Additionally, we need to add in any necessary tests.

I might enlist some help on this cleanup in winter 2022. If so, I'll provide some more detailed instructions.

Basic workflow:

1. checkout bring-some-of-3la-pldi-push-main-to-main
2. choose a file from the list below and run git diff bring-some-of-3la-pldi-push-main-to-main..main
3. inspect the changes and determine which are appropriate to bring to main
4. run git checkout --patch main -- <filename>. This will allow you to revert changes hunk-by-hunk. For the changes which shouldn't be merged, hit y to tell git to create a revert change for the change. For the changes which should be merged, hit n -- no revert change will be created.
5. once you finish the checkout, commit all of the reverts for that file. They should already be staged.

Each major change being brought to main should be tested.

This will likely also result in changes to TVM, which will hopefully get merged into mainline. We have a similar 3la-pldi-push-main branch on our TVM repo, for which we'll need to do this same process.

List of files:

• [X] .github/workflows/build-and-test.yml
• [x] .github/workflows/rustfmt.yml
• [ ] Cargo.toml
• [x] Dockerfile
• [x] examples/glenside-cli.rs
• [x] models/lstm-for-pldi-pattern.relay
• [x] models/lstm-for-pldi.relay
• [x] models/resmlp.relay
• [x] models/resnet.relay
• [x] models/transformer.relay
• [x] src/codegen.rs
• [x] src/extraction/ilp.rs
• [ ] src/extraction/mod.rs
• [ ] src/language/from_relay/mod.rs
• [ ] src/language/interpreter.rs
• [ ] src/language/language.rs
• [ ] src/language/rewrites.rs
• [ ] tests/3la-glenside.rs
• [ ] tests/3la-resnet.rs
• [ ] tests/codegen-mlp.rs
• [ ] tests/conv1d-flexmatch.rs
• [ ] tests/lstm_relay_to_glenside.rs
• [ ] tests/mobilenet-relay-to-glenside.rs
• [ ] tests/mobilenet-try-to-run-rewrites.rs
• [ ] tests/resmlp-linear-rewrites.rs
• [ ] tests/resnet18_relay_to_glenside.rs
• [ ] tests/tensorize-conv2d-with-padding-and-splitting-from-command-line.rs
• [ ] tests/tensorize-conv2d-with-padding-and-splitting-with-blocking-from-command-line.rs
• [ ] tests/tensorize-conv2d-with-padding-and-splitting.rs
• [ ] tests/transformer.rs

specific todos:

• [x] relay_op_softmax test in codegen shouldn't be ignored
• [x] from_relay() needs to be updated so that it doesn't return a tuple
• [x] Compilation needs to use the new relay->glenside rewrites
• [x] Fix conv2d NHWC typechecking in language.rs (see warn!)
• [x] conv2d from_relay should respect use_opaque_operators_for. Currently, the glenside impl is always generated.
• [x] We can't point to 3la-tvm, should point to mainline TVM if possible
• [x] Update to new egg

[gussmith23]

I think the addition of name_to_dtype should be kept. We can perhaps reimplement it if needed, but it's probably right to have it, esp. for compiling back to Relay.

[gussmith23]

I think Mike's changes to ILP are all good, but we should make sure all the tests pass.

[gussmith23]

We no longer manually union things using Mike's approach from the PLDI push. The relay->glenside rewrites should handle everything correctly now.

[gussmith23]

I think we should remove relay_shape. The reason it's there is so that we can let access pattern shapes take precedence in merge(), as far as I can tell. That is, if an eclass only has Relay operators, then its access patterns will look fairly arbitrary. (Sometimes i put all the dimensions in shape, sometimes I put them all in item_shape...)

Two cleaner solutions:

• Short term: replace relay_shape with contains_only_relay_nodes. If this bool is true, then there are only Relay nodes in the eclass and thus we know to treat the access pattern shape as a tensor shape (i.e. concatenate shape + item_shape).
• Long term: decide on a cleaner way to represent Relay shapes in terms of access patterns.

[gussmith23]

I'm rewriting the Analysis for Glenside to accomodate for Relay. Previously, Mike did this by including a "relay shape" field, with some additional logic in merge() to handle the cases where a node might have a relay shape and no access pattern shape or vice versa.

I think we can simplify this to just have a flag which keeps track of whether an eclass contains only Relay nodes. Then, if an eclass contains only Relay nodes, we interpret its access pattern shape as a tensor shape (i.e. if its access pattern shape is ((a,b),(c,d)), we interpret it instead as just (a,b,c,d)). We expect that the underlying tensor shapes must match, but that we may have to take the access pattern shape from the eclass with "real" glenside nodes, as that should have the "correct" access pattern data.

All of this is a consequence of not having a disciplined way of defining relay shapes in terms of access pattern shapes. I.e. we could always put them in the shape field, or always put them in the item shape field, but I do a mix of both.

[gussmith23]

One of the main problems with this currently is the fact that something like the following can happen:

Eclass A starts with a relay node. It points to eclass B, which starts with a Relay node. Eclass A gets rewritten and now includes a Glenside node. Its access pattern data is recalculated and now considered "correct". Eclass B gets rewritten and now includes a Glenside node. Its access pattern data is also recalculated. Assume the access pattern shape changes, and this new version is now considered "correct". This rewriting triggers the recalculation of analysis data for eclass A (see process_unions in egg). Because the eclass B's shape changes, eclass A's shape changes. We try to merge that new analysis data with the old analysis data for eclass A, and it fails, because both are "correct" but they differ.

The thing we're missing here is that, just because a node is a Glenside node, it doesn't mean its access pattern data is correct/unchanging/"settled". Instead, we have to know that all of its children have settled their shapes as well. How do we know if a node's shape is settled? I think it's just a matter of whether

1. the node is a Glenside node and not a Relay node (which is what we currently track)
2. all of the node's children are settled (which we currently don't track).

So I think it's just a matter of adding that second piece of information. When we build an analysis for a new node,

• If it's a Relay node, "settled" is false
• If it's a Glenside node, "settled" = &(settled of all access-pattern-typed children).

[gussmith23]

Okay, I've resolved the above. But still have some residual errors w/ access patterns not calculating correctly.

[gussmith23]

Current issue: using access-shape to make a shape literal doesn't work when that shape literal might need to change b/c the underlying shape isn't yet finalized. The solution is to use get-access-shape to point to the shape of an existing eclass...except when that's not actually possible! I.e. when the shape isn't the shape of another eclass.

All of this is making me think I should just rewrite the damn eclass analysis once and for all! We would just put all elements in a vec and have an optional usize for the access axis.