[PATCH] D133166: [mlir][MemRef] Canonicalize extract_strided_metadata(subview)

[Nicolas Vasilache via Phabricator via llvm-commits]

nicolasvasilache added a comment.

Thanks for making progress on this @qcolombet !

Here is a first round of comments to reduce the impl size and make it more idiomatic.



================
Comment at: mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp:1222
+      Value subStride;
+      if (subview.isDynamicStride(i))
+        subStride = subview.getDynamicStride(i);
----------------
We should have something in StaticValueUtils.h|cpp to hide away the if/else complexity and return an OpFoldResult.
Could you reuse or extend ?


================
Comment at: mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp:1228
+
+      Value newStride = rewriter.createOrFold<arith::MulIOp>(
+          origLoc, newExtractStridedMetadata.getStrides()[i], subStride);
----------------
Can we use a single AffineApplyOp here and below?
At the higher levels of abstraction this is a better way of representing such indexing computations (i.e. 1 op instead of N ops).

See e.g. https://sourcegraph.com/github.com/llvm/llvm-project/-/blob/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp?L154
You can also make the variadic process of defining symbols bound to a position better by using something like: https://sourcegraph.com/github.com/llvm/llvm-project/-/blob/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp?L281


================
Comment at: mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp:1255
+          rewriter.createOrFold<arith::MulIOp>(origLoc, subOffset, strides[i]);
+      offset = rewriter.createOrFold<arith::AddIOp>(origLoc, offset,
+                                                    scaledSubOffset);
----------------
Same above, a lot of this could be compressed down to almost a one liner


================
Comment at: mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp:1271
+    // the subview.
+    for (int i = 0; i < sourceRank; ++i) {
+      if (droppedDims.test(i))
----------------
Can this be factored out into a helper function that would live next to e.g. `SubViewOp::inferResultType` ?
It could return a vector of OpFoldResult to avoid passing a builder and then StaticValueUtils can be used or extended to materialize the constants here.

All this should then become much more idiomatic (and shorter).


================
Comment at: mlir/test/Dialect/MemRef/canonicalize.mlir:868
+//
+//   CHECK-DAG: %[[DIM0_OFFSET:.*]] = arith.muli %[[STRIDES]]#0, %[[C3]]
+//   CHECK-DAG: %[[DIM1_OFFSET:.*]] = arith.muli %[[STRIDES]]#1, %[[C4]]
----------------
with the suggestion above you'd only check 1 `affine.apply` op here with the properly captured AffineMap.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D133166/new/

https://reviews.llvm.org/D133166