Remove rank 0 affine.parallel when canonicalizing

[tzerrell]

AffineParallelOps of rank 0 (i.e., those with no induction variables) are superfluous. This adds a canonicalizer for AffineParallelOps that removes them if they are rank 0.

Submitting here for verification that this correctly enables the changes we want on the PlaidML side before submitting upstream. We can also land this and point PlaidML here if we decide we want to move faster than the LLVM/MLIR review process. We probably shouldn't land until we've rebased this on top of #5 and landed plaidml/plaidml#1062 on the PlaidML side, as there are some minor conflicts that will need to be resolved.

[flaub]

Hmm, the problem with this is how do we deal with kernel outlining? Wouldn't we want to not remove affine.parallel with no induction variables?

[tzerrell]

@jbruestle what do you think? One option I suppose is to not canonicalize if we want that use case -- although deliberately avoiding canonicalization seems a bit sketchy and also perhaps we'll need to canonicalize for other reasons? Another possibility that occurs to me is some other wrapper to represent kernel outlining. I don't fully understand the requirements and goals of kernel outlining though.

We can always just use the code in plaidml/plaidml#1054, which works just fine. This is here because having a no-induction-variables affine.parallel feels non-canonical and like something that should be simplified away (and doing so would address the same problems as that PR). But if we have need for a rank 0 affine.parallel then maybe we shouldn't be removing it

[flaub]

I think it's true that while semantically, unwinding an empty affine.parallel is equivalent, I wonder if we should be considering the structure of affine.parallel (like how many nests there are) as some other bit of information for other passes to use. For instance I can imagine that some pass might want to think of the outermost affine.parallel as the loop over threads, and then another affine.parallel to loop over something else. Even if the iteration domain is 1, which can either be represented as affine.parallel () = () to () or affine.parallel (i) = (0) to (0), they both imply some structure perhaps as just a way to group a collection of ops together.

Then again, maybe it's not a good idea to use affine.parallel like that (as a means of grouping ops). I've often thought maybe we should have a kernel op to represent this grouping.

Basically, in stripe, we had assumed that the outermost block was always the 'kernel' level. How do we represent such a thing in affine dialect, or do we even need such a thing?

[jbruestle]

So for one thing, a kernel that operates on one element is slow inside a affine.parallel or outside of it, and so having kernel outlining rebuild an empty 1 element kernel when it doesn't find an affine.parallel seems very reasonable. On the other hand, that's not the only place where 'structuring' the program in a platform specific way makes sense, for example on the GPU we have 1) 'grid/group' affine.parallel 2) thread affine.parallel.

What about this: we don't allow canonicalization if the affine.parallel has any attributes on it? Or maybe we look for a specific 'unit' attribute like 'is_structural' or whatever?

It seems to me like 1) The general case of canonicalization of empty parallels is good 2) 'Structuring' affine.parallel like constructs will exist at some point, and making them each a new op is annoying causes all kinds of issues, so a compromise is perhaps the best approach. Notably, the attribute trick was what I did for the Stripe dialect, which did have an empty PF canonicalization.

[jbruestle]

However, it's also reasonable for GPU lowering to look for a 'grid' AP, and a 'thread' AP, and if either is missing, add an empty loop.

[jbruestle]

The first option (limiting canonicalization when an attribute it present) in some sense pushes our design upstream more, and the second option (re-introducing affine.parallels when a specific one isn't present) puts a bit more work on the lowering, but is more 'neutral'. I don't have a strong opinion on either side, but if we are going to go with the attribute to prevent canonicalization, we will need to describe why it's there when we upstream.

[flaub]

I have a feeling that upstream won't accept 'special' attributes as a way to control canonicalization. My guess is that the complaint will be that attributes should not 'leak' or otherwise have semantics outside of the dialect they were defined in. I don't know of any other situation where attributes/markers are used to adjust canonicalization.

Personally I find it a pragmatic solution, but then I wonder what the attribute should be named, and if it should be defined on the affine.parallel as an official attribute rather than an ad-hoc one (probably official).

[tzerrell]

That makes sense to me. Since attributes in general might have structural information I will for now at least draft a version that only canonicalizes when there are no unexpected attributes. Bah, didn't see the latest posts