Fold memref.expand_shape and memref.collapse_shape ops into
their memref/affine load/store ops.
Details
Diff Detail
- Repository
- rG LLVM Github Monorepo
Event Timeline
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. 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. | |
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. |
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. |
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. |
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. | |
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. |
@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 ? |
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.
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. |
@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.
mlir/include/mlir/Dialect/Utils/IndexingUtils.h | ||
---|---|---|
18 | You don't need this include. Please use fwd decl. |
utils/arcanist/clang-format.sh | ||
---|---|---|
54 | please revert this change |
@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. |
A few minor things.
- Please add a commit summary; it's currently blank.
- Drop the full stop at the end of the commit title.
- 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. |
This comment is outdated.