This is an archive of the discontinued LLVM Phabricator instance.

Fold memref.expand_shape and memref.collapse_shape ops
ClosedPublic

Authored by arnab-oss on Jul 1 2022, 5:52 AM.

Details

Summary

Fold memref.expand_shape and memref.collapse_shape ops into
their memref/affine load/store ops.

Diff Detail

Event Timeline

arnab-oss created this revision.Jul 1 2022, 5:52 AM
Herald added a project: Restricted Project. · View Herald TranscriptJul 1 2022, 5:52 AM
arnab-oss requested review of this revision.Jul 1 2022, 5:52 AM
arnab-oss updated this revision to Diff 441671.Jul 1 2022, 6:01 AM

Discarding unnecessary changes.

Thanks for improving this part of the codebase @arnab-oss!
Here is a first batch of comments.

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
31

nit: typo: "of a load/store" here and below

56

Please generalize this a bit and add as a helper in include/mlir/Dialect/Utils/IndexingUtils.h, this will be useful in other places too.

I am surprised there isn't already something like this, I only see utilities to linearize / delinearize atm.

67

This feels like a linearization similar to the one in IndexingUtils.h but more generally applicable.
Could you try generalizing the AffineExpr building part into a new IndexingUtils.h and compose that with the extra reorderings?

The more it looks like:

AffineExpr srcIndexExpr = linearAffineExpression(suffixProduct);
SmallVector<Value> dynamicIndices = gather(indices, groups);
sourceIndices.push_back(rewriter.create<AffineApplyOp>(...));

the better the composability and overall code idiom power.

108

Same comments as above but for a delinearized form.

178

Have not yet looked at this deeply but please consider a similar type of refactoring / reuse as proposed above if possible.
I'll come back to it once the first batch of comments is resolved.
(I see you are merely moving code here but if there are opportunities to improve I'd take them).

245

Looking at the amount of code below I am wondering if you are going overboard with templates.

Seems it would be quite fewer code to have 3 separate LoadOpOfXXXFolder and inline the impl at the right place.
What you save in not typing the pattern decl and the matchAndRewrite decl you pay at least 4-5x the price in the extra indirection and templating logic.

This would be a great enhancement for code generation! -- to be able to de-abstract these ops out if/when needed. Would have been great to have these from the start!

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
2

Update summary line: memref.subview -> memref aliasing ops.

This would be a great enhancement for code generation! -- to be able to de-abstract these ops out if/when needed. Would have been great to have these from the start!

Feel free to contribute when you need something ! :)

arnab-oss updated this revision to Diff 444272.Jul 13 2022, 8:27 AM
arnab-oss marked 4 inline comments as done.

Addressed review comments.

arnab-oss marked an inline comment as done.

Addressd review comments.

arnab-oss marked an inline comment as done.Jul 17 2022, 11:34 PM

ping @nicolasvasilache can you please re-review?

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
178

I feel this part of the code usage is specific to this case, and is not beneficial to expose it as a utility function.

It looks like the commit summary has got into the title, and the former is missing!

arnab-oss retitled this revision from Fold memref.expand_shape and memref.collapse_shape ops, when they are the parent op of source value of memref/affine load/store ops. to Fold memref.expand_shape and memref.collapse_shape ops..Jul 18 2022, 1:03 AM
mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
123

This seems super artificial, I would rather duplicate this affine builder

void AffineLoadOp::build(OpBuilder &builder, OperationState &result,
                         Value memref, ValueRange indices) {

to allow taking ArrayRef<OpFoldResult> and just pass fill your sourceIndices with b.getIndexAttr(0).

178

faire enough, thanks for looking into it.

270

This decl, its def can be folded in all its uses now I think.
This would be the part that reduces the code size.

496

I think you can now:

llvm::TypeSwitch<Operation *, void>(loadOp)
        .Case<affine::LoadOp, memref::LoadOp>([&](auto op) {
          rewriter.replaceOpWithNewOp<decltype(op)>(
            loadOp, expandShapeOp.getViewSource(), sourceIndices);
        })
        .Default([](Operation *) { ;});

And drop all the overloads and decls?

mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
278

Please also add some tests for memref.load/store and vector.transfer.

arnab-oss marked 2 inline comments as done.Jul 25 2022, 3:26 AM

@nicolasvasilache can you please take a look at my comments?

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
123

Hi, I'm sorry I do not get your comment. Can you please elaborate?

270

I think I cannot unify all the cases involving different varieties of load/store ops, as their build functions expect different kinds and numbers of arguments.

496

I think this is not possible. Please take a look at my comment above.

Sorry was OOO the last 2 weeks, replied inline.

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
496

This is precisely why I used a specific TypeSwitch + auto + decltype to resolve the type inside the Case lambda. To spell it more, you should be able to:

llvm::TypeSwitch<Operation *, void>(loadOp)
        .Case<affine::LoadOp, memref::LoadOp>([&](auto op) {
          rewriter.replaceOpWithNewOp<decltype(op)>(
            loadOp, expandShapeOp.getViewSource(), sourceIndices);
        })
        .Case<vector::TransferReadOp>([&](auto transferReadOp) {
          if (transferReadOp.getTransferRank() == 0) {
            // TODO: Propagate the error.
            return;
          }
          rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
              transferReadOp, transferReadOp.getVectorType(), subViewOp.source(),
              sourceIndices,
              getPermutationMapAttr(rewriter.getContext(), subViewOp,
                                    transferReadOp.getPermutationMap()),
              transferReadOp.getPadding(),
              /*mask=*/Value(), transferReadOp.getInBoundsAttr());
        })
        .Default([](Operation *) { ;});

Did you try and hit compilation errors? If so, could you paste the error message ?

arnab-oss added a comment.EditedAug 12 2022, 4:24 AM

I'm getting the following error:

/home2/arnab/llvm-project/mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp:447:30: error: ‘class mlir::vector::TransferReadOp’ has no member named ‘getValue’
  447 |             storeOp, storeOp.getValue(), subViewOp.source(), sourceIndices);

AffineStoreOp and memref::StoreOp have the getValue(), but vector::TransferReadOp does not.

I'm getting the following error:

/home2/arnab/llvm-project/mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp:447:30: error: ‘class mlir::vector::TransferReadOp’ has no member named ‘getValue’
  447 |             storeOp, storeOp.getValue(), subViewOp.source(), sourceIndices);

AffineStoreOp and memref::StoreOp have the getValue(), but vector::TransferReadOp does not.

Sorry I missed this.. This seems like the type of errors that the TypeSwitch solution I suggested should easily handle .. Can you post another WIP PR which shows how you are trying to do this?
In any case, adding a getValue to TransferReadOp or renaming the result vector to value is also an easy normalization of APIs.

If you find this really difficult, please just add more tests as requested and feel free to land as is.
I can pick up the slack in a followup commit.

bondhugula accepted this revision.Aug 19 2022, 6:51 PM

I'm getting the following error:

/home2/arnab/llvm-project/mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp:447:30: error: ‘class mlir::vector::TransferReadOp’ has no member named ‘getValue’
  447 |             storeOp, storeOp.getValue(), subViewOp.source(), sourceIndices);

AffineStoreOp and memref::StoreOp have the getValue(), but vector::TransferReadOp does not.

You can just make the APIs (accessor names here) consistent.

This revision is now accepted and ready to land.Aug 19 2022, 6:51 PM
bondhugula added inline comments.Aug 19 2022, 6:52 PM
mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h
41–42

This comment is outdated.

96

This one as well. Please update the doc comments to reflect the generalization.

bondhugula requested changes to this revision.Aug 19 2022, 6:55 PM

There are still several style and efficiency issues (unnecessary copies).

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
50–51

This would lead to unnecessary copies - please don't use SmallVector's by copy this way; groups is only read below. Normally const ref, but ArrayRef here.

97

Likewise.

This revision now requires changes to proceed.Aug 19 2022, 6:55 PM
bondhugula accepted this revision.Aug 19 2022, 6:59 PM
bondhugula added inline comments.
mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
50–51

Sorry, looks like I didn't see the RHS here. This is a SmallVector returned from a function. So this is correct as is.

97

Nothing to change.

This revision is now accepted and ready to land.Aug 19 2022, 6:59 PM
arnab-oss updated this revision to Diff 454442.Aug 22 2022, 3:39 AM
arnab-oss marked 9 inline comments as done.

Addressed comments

@nicolasvasilache I've made the required changes in the API. Also added tests featuring memref.load/store. However I've not handled folding related to vector transfer ops, so no tests involving those ops.

bondhugula accepted this revision.Aug 22 2022, 5:28 AM
bondhugula added inline comments.
mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
1488–1490

Doc comment here. You can mention that this is added for uniformity with other similar ops.

1488–1490

Nit: This will fit in a single line.

bondhugula added inline comments.Aug 22 2022, 5:31 AM
mlir/include/mlir/Dialect/Utils/IndexingUtils.h
18

You don't need this include. Please use fwd decl.

arnab-oss updated this revision to Diff 454473.Aug 22 2022, 5:53 AM
arnab-oss marked 3 inline comments as done.

Addressed comments.

utils/arcanist/clang-format.sh
54

please revert this change

arnab-oss marked an inline comment as done.

Removed wrong changes in clang-format.sh

nicolasvasilache accepted this revision.Aug 23 2022, 2:22 AM

@arnab-oss thanks for refactoring to a nicer outcome.

Please address the last 2 points I raised and let's ship this!

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
344

Please add an llvm_unreachable here and all other places below, we don't want to silently fail if/when new ops get added in the future.

mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp
50

Seems unrelated to this PR.
Even if it has no effects on the logic, I'd rather keep separate pieces separate.

arnab-oss updated this revision to Diff 454790.Aug 23 2022, 4:43 AM
arnab-oss marked an inline comment as done.

Addressed comments by nicolas.

arnab-oss updated this revision to Diff 454794.Aug 23 2022, 4:46 AM

Minor changes.

bondhugula accepted this revision.Aug 25 2022, 11:56 PM

A few minor things.

  1. Please add a commit summary; it's currently blank.
  2. Drop the full stop at the end of the commit title.
  3. There is a whitespace issue in the test case file: $ git diff --check HEAD~.
mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
418–419

Indent by two to be consistent.

arnab-oss updated this revision to Diff 455821.Aug 26 2022, 1:14 AM

Minor NFC changes.

arnab-oss edited the summary of this revision. (Show Details)Aug 26 2022, 1:29 AM
arnab-oss edited the summary of this revision. (Show Details)Aug 26 2022, 11:16 PM
bondhugula retitled this revision from Fold memref.expand_shape and memref.collapse_shape ops. to Fold memref.expand_shape and memref.collapse_shape ops.Aug 27 2022, 6:27 PM
bondhugula edited the summary of this revision. (Show Details)
This revision was automatically updated to reflect the committed changes.
mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp