Page MenuHomePhabricator

Fold memref.expand_shape and memref.collapse_shape ops.
Needs ReviewPublic

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

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.Sun, Jul 17, 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..Mon, Jul 18, 1:03 AM
mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
122

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.

269

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

495

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.Mon, Jul 25, 3:26 AM

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

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

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

269

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.

495

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
495

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.EditedFri, Aug 12, 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.