Support alignment of stores on specific dimensions only

[manopapad]

Requested by @elliottslaughter. Copying offline discussion:

Elliott Slaughter: So my partitions aren't going to be quite this simple. I have a set of arrays with shapes that look like (A, X), (X, B), (C, X, D) where A, B, C, D and X are all different variables, and I want to partition only on X.

Manolis: So what you would do today would be to make all arrays the same shape. E.g. to align a: (A, X) and b: (C, X, D) you'd broadcast a: (A, X) to a': (A, X, C, D) and b: (C, X, D) to b': (A, C, X, D), then need to transpose b': (A, C, X, D) to b'': (A, X, C, D). Then a' and b'' have the same shape, so you can align them when you pass them to the task. You have all these useless dimensions, they're not causing copies, but you do need to ignore them in the task body.

Elliott: I can test like this but it would really be better if this gets fixed properly before I have to show it to users

[elliottslaughter]

This is needed for the SLAC/LANL project.

It turns out that in the real code, there are multiple dimensions you can partition along. To expand my example, it might look something like:

• Array 1 with dimensions (A, X)
• Array 2 with dimensions (X, Y)
• Array 3 with dimensions (Y, X, D)

And then I'd write a constraint like:

align(("array1", 1), ("array2", 0), ("array3", 1)),
align(("array2", 1), ("array3", 0))

The point being that we can down slice the X or the Y, or both.