Page MenuHomePhabricator

[mlir] Use affine.apply when distributing to processors
ClosedPublic

Authored by antiagainst on Mar 8 2021, 4:47 AM.

Details

Summary

This makes it easy to compose the distribution computation with
other affine computations.

Diff Detail

Event Timeline

antiagainst created this revision.Mar 8 2021, 4:47 AM
antiagainst requested review of this revision.Mar 8 2021, 4:47 AM
mlir/lib/Transforms/Utils/LoopUtils.cpp
2160

Thanks!
Note that we discourage chains of affine.apply operations are discouraged.
As an alternative you could call this iteratively:

AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map,
                                      ArrayRef<Value> operands);

Even better, just iterate to create the proper AffineExpr so that you never insert ops that would become dead in the first place.
Then you can just call:

void canonicalizeMapAndOperands(AffineMap *map, SmallVectorImpl<Value> *operands);

to cleanup the expressions and operands and then only create the AffineApplyOp.

bondhugula added inline comments.
mlir/lib/Transforms/Utils/LoopUtils.cpp
2160

Yes, but in IR building instances like these, using chains of affine.apply I think makes the code much cleaner and manageable. The costs of canonicalizing these are nothing when consider everything else that's happening. I don't think one should try too hard to make composed things at the expense of repeating the code that performs such canonicalization all over the place. Note that composing an affine apply + canonicalizeMapAndOperands doesn't still give you the fixed point simplified form. You need simplifyAffineMap as well in the mix. -canonicalize takes care of it.

The place I would recommend not creating such affine.apply ops is where they can be composed/pulled/folded into the op that's consuming the result (like load/store/affine.for).

If you really want the simplified form free of these chains, just run the local canonicalizer (applyPatternsAndFoldLocally).

mravishankar accepted this revision.Mar 8 2021, 3:21 PM

Awesome! Thanks Lei for tracking this down! LGTM

mlir/lib/Transforms/Utils/LoopUtils.cpp
2160

In this case though it is (a*b + c). I think just creating that map and having a single affine.apply doesnt hurt readability. I'd argue it increases it cause you know exactly the final expression you are generating.

This revision is now accepted and ready to land.Mar 8 2021, 3:21 PM
antiagainst added inline comments.Mar 9 2021, 5:35 AM
mlir/lib/Transforms/Utils/LoopUtils.cpp
2160

I agree with @bondhugula. I don't see the needs to create well canonicalized IRs. That should be the responsibility of dedicated passes. It's a better separation of concerns to me and good example of pattern/pass composability.