Operator and access fusion

[gussmith23]

For 3la work, it seems like it'd be useful to be able to represent fused operators in Glenside.

Something like

(compute add
 (pair
  (compute dot product
   (cartprod ?a ?b))
  ?c))

would be nice to represent as

(compute-fused (list dot-prod pair add)
  (cartprod ?a ?b) ?c))

I think this is doable.

First step would be to represent access operators as we represent compute operators, i.e. (access op args). This allows us to make the "fusion list" that I showed above, (list dot-prod pair add), of access pattern ops and compute ops.

Type checking would then go through each compute/access op, performing type checking on them in sequence. There will be some limitations early on, I think, with how type checking can proceed. To illustrate, let me step through how we'd type-check/do shape inference on the above example.

• We would iterate through the list of compute/access ops and shape-infer/type-check them in sequence. We start first with dot-prod. Dot prod takes one argument, so we pop the first argument off the list of args, in this case the cart-prod expression.
  • Here we hit our first limitation: if we implement it the above way, then we need to know how many args each op takes. Alternatively, we represent the list of arguments as a list of lists.
  • Another big limitation: I'm not sure how to represent any access patterns that require non-tensor arguments. One way would be to differentiate between data args and attributes, as TVM/Relay does. So for example, the current access operator, which we might rename to at so that it's (access at ...), would take one attribute (the access axis) and one data argument (the tensor to access). Then, to represent it in a fusion pipeline, we'd have (list ... (at 2) ...), where in this case the operator is constructed with its attributes.
• To type-check, we would just run the same cartprod type checking routine we already have in Glenside.
• Then we'd store the resulting intermediate type and move on to pair. Pair takes two args, and we have one provided by the intermediate value, so we pop another off the list.
  • Here's another limitation: we have to make an assumption about argument order. so perhaps we assume that the intermediate value is always the first arg. In this case, it doesn't matter, but in other cases it would. This could be fixed by replacing the list with a full AST structure, showing where arguments lie.
• We'd run type checking on the pair, to produce another intermediate type.
• Lastly, add takes only a single arg, which is our intermediate. So we just run that through type checking for add and we're done.

I think this works. We could do this quickly by doing the following:

• Introduce a new apply-access-op construct: (apply-access-op <access-op> <arg>...)
• Introduce new access ops for the accesses we need: in the example above, we only need pair.
• Introduce compute-fused and implement the type checking routine I described above.

The big question is, does this work for Akash?