Page MenuHomePhabricator

[mlir]: Add folding for DimOp with AllocaOp operand
ClosedPublic

Authored by akuegel on Feb 26 2021, 4:42 AM.

Details

Summary

This adds AllocaOp to the operand ops that we support for folding.
This folding can be very useful for extent tensor computations.

Diff Detail

Event Timeline

akuegel created this revision.Feb 26 2021, 4:42 AM
akuegel requested review of this revision.Feb 26 2021, 4:42 AM
bkramer added inline comments.Feb 26 2021, 6:35 AM
mlir/lib/Dialect/StandardOps/IR/Ops.cpp
1528 ↗(On Diff #326655)

This should also be safe for any AllocLikeOp

1537 ↗(On Diff #326655)

I don't understand why this must be a rank op

akuegel updated this revision to Diff 327034.Mar 1 2021, 12:03 AM

Make the pattern more general.

akuegel marked an inline comment as done.Mar 1 2021, 12:05 AM
akuegel added inline comments.
mlir/lib/Dialect/StandardOps/IR/Ops.cpp
1537 ↗(On Diff #326655)

You are completely right, it doesn't have to be a rank op. This was my motivating example, but no need to restrict it to that.

akuegel retitled this revision from [mlir]: Add canonicalization for dim of 1D alloc of size rank. to [mlir]: Add canonicalization for dim of 1D AllocLikeOp.Mar 1 2021, 12:06 AM
bondhugula requested changes to this revision.Mar 1 2021, 12:56 AM
bondhugula added a subscriber: bondhugula.
bondhugula added inline comments.
mlir/test/Dialect/Standard/canonicalize.mlir
137–146

This folding already exists and happens upstream on this test case. Is this an (extra) test case for existing functionality?

The fold pattern on DimOp already handles this for any dimensionality memrefs.

155

Please change the name of this function - it gives the impression that this has something to do with the RankOp but in reality it has nothing to. It could as well be the result of any operation or even a block argument. Also, this can be done for any dimensional memref.

157

This is really a folding pattern and should be implemented as part of DimOp::fold - it allows faster progress in the canonicalizer (greedy rewrite driver) and that's the standard approach. All such replacements on other ops are actually implemented on the folding hook - some that were canonicalization patterns earlier were migrated almost a year ago to the folding patterns. The fold methods are also called during the greedy rewrite driver. It just allows more and earlier simplification with this being part of fold.

This revision now requires changes to proceed.Mar 1 2021, 12:56 AM
akuegel updated this revision to Diff 327051.Mar 1 2021, 1:25 AM

Move canonicalization pattern to folding.

akuegel marked an inline comment as done.Mar 1 2021, 1:29 AM
akuegel added inline comments.
mlir/test/Dialect/Standard/canonicalize.mlir
137–146

You are right, I just hadn't seen that this already works. I started with noticing that dim(alloca) of a 1D memref was not canonicalized, and that was what I initially wanted to fix. I removed the test case again.

157

Thanks, that makes sense :)
I moved it to fold.

akuegel retitled this revision from [mlir]: Add canonicalization for dim of 1D AllocLikeOp to [mlir]: Add folding for dim of operand with rank 1.Mar 1 2021, 1:31 AM
akuegel edited the summary of this revision. (Show Details)

I think I did the requested changes. Can you please take another look? :)

A couple more comments @akuegel

mlir/lib/Dialect/StandardOps/IR/Ops.cpp
1507–1514 ↗(On Diff #327051)

Can these two ifs be combined into a single one with the dyn_cast_or_null on AllocLikeOp?

mlir/test/Dialect/Standard/canonicalize.mlir
147

Could you add extra lines to test the case when the alloc operand is a block argument? That should just fold as well.

bondhugula requested changes to this revision.Mar 2 2021, 2:14 AM
bondhugula added inline comments.
mlir/lib/Dialect/StandardOps/IR/Ops.cpp
1446–1447 ↗(On Diff #327051)

I still don't understand what the connection with rank 1 is. This folding for DimOp or the propagation of the AllocLikeOp operand can happen for any ranked memref operand of a DimOp.

This revision now requires changes to proceed.Mar 2 2021, 2:14 AM
bondhugula added inline comments.Mar 2 2021, 2:27 AM
mlir/lib/Dialect/StandardOps/IR/Ops.cpp
1447 ↗(On Diff #327051)

if the code is correct <- This would be a bit misleading. I think you are referring to undefined behavior which would happen if the index of a dim op is >= rank of memref/tensor. Basically, you would be exploiting undefined behavior.

if the code is correct -> for the behavior to be defined ?

@mehdi_amini @rriddle Should we exploiting undefined behavior this way?

1449–1454 ↗(On Diff #327051)

I think I now see what you were handling and the other thing I was thinking about. I was just thinking (separately from the intention of your patch) that you could get the indexth operand from the AllocLikeOp that's defining this memref and fold the DimOp by replacing its result with that value. A Value is returned as OpFoldResult in such cases. It is orthogonal to what you are doing here.

akuegel updated this revision to Diff 327699.Mar 3 2021, 12:47 AM

Address review comments.

akuegel added inline comments.Mar 3 2021, 12:48 AM
mlir/lib/Dialect/StandardOps/IR/Ops.cpp
1447 ↗(On Diff #327051)

I thought about this a bit more, and even if we allow exploiting undefined behavior (which seems to be not clear), it is not worth it. If the emitted code knows it is dealing with a 1D memref, I think the index operand of DimOp would always be a constant. So I removed this part again.

1449–1454 ↗(On Diff #327051)

If it is orthogonal, I guess it could be done in a separate patch.
For now, I only care about that the folding also works for AllocaOp :)

1507–1514 ↗(On Diff #327051)

As far as I can tell, the class AllocLikeOp only exists in the Ops.td file, it is not generated as part of the code created with tablegen. If I look at the generated AllocOp class, it derives from ::mlir::Op, not from AllocLikeOp. But I am not too familiar with the MLIR infrastructure yet, so it is possible I overlooked something.
I have checked for instances dealing with both AllocOp and AllocaOp, they seem to always use a template parameter AllocLikeOp. But I guess we cannot use a template for the fold method?

akuegel retitled this revision from [mlir]: Add folding for dim of operand with rank 1 to [mlir]: Add folding for DimOp with AllocaOp operand.Mar 3 2021, 12:50 AM
akuegel edited the summary of this revision. (Show Details)
akuegel marked an inline comment as done.

Please take another look :)
I have removed the controversial part of the change.

bondhugula accepted this revision.Mar 13 2021, 2:40 AM

LGTM.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp
1507–1514 ↗(On Diff #327051)

Does it compile when you use AllocLikeOp? Yes, it only appears in the td file. A class doesn't have to inherit from another for cast/dyn_cast etc. to work. (For eg. Operation * cast to any derived OpTy.)

This revision is now accepted and ready to land.Mar 13 2021, 2:40 AM
akuegel updated this revision to Diff 330912.Mar 16 2021, 2:30 AM

Fix test comment.

Does it compile when you use AllocLikeOp? Yes, it only appears in the td file. A class doesn't have to inherit from another for cast/dyn_cast etc. to work. (For eg. Operation * cast to any derived OpTy.)

Unfortunately it doesn't compile. No C++ class seems to be generated for it (I didn't find it in the Ops.h.inc (now MemRefOps.h.inc) file where e.g. AllocOp and AllocaOp classes are forward declared). Maybe it would be generated in some other file, and I would need to include that other file?

This revision was landed with ongoing or failed builds.Mar 16 2021, 2:39 AM
This revision was automatically updated to reflect the committed changes.
Harbormaster completed remote builds in B94000: Diff 330912.