gussmith23 / glenside

A pure, low-level tensor program representation enabling tensor program optimization via program rewriting. See the web demo at https://gussmith23.github.io/glenside-web-demo/
71 stars 10 forks source link

Operator and access fusion #176

Open gussmith23 opened 2 years ago

gussmith23 commented 2 years ago

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.

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

The big question is, does this work for Akash?