[Mlir-commits] [mlir] [MLIR] Fix incorrect memref::DimOp canonicalization, add tensor::DimOp canonicalization (PR #84225)

[Mehdi Amini]

================
@@ -1069,6 +1069,52 @@ OpFoldResult DimOp::fold(FoldAdaptor adaptor) {
   return {};
 }
 
+namespace {
+/// Fold dim of a memref reshape operation to a load into the reshape's shape
+/// operand.
+struct DimOfMemRefReshape : public OpRewritePattern<DimOp> {
+  using OpRewritePattern<DimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DimOp dim,
+                                PatternRewriter &rewriter) const override {
+    auto reshape = dim.getSource().getDefiningOp<ReshapeOp>();
+
+    if (!reshape)
+      return rewriter.notifyMatchFailure(
+          dim, "Dim op is not defined by a reshape op.");
+
+    if (dim.getIndex().getParentBlock() == reshape->getBlock()) {
+      if (auto *definingOp = dim.getIndex().getDefiningOp()) {
+        if (reshape->isBeforeInBlock(definingOp))
+          return rewriter.notifyMatchFailure(
+              dim,
+              "dim.getIndex is not defined before reshape in the same block.");
+      } // else dim.getIndex is a block argument to reshape->getBlock
+    } else if (!dim.getIndex().getParentRegion()->isProperAncestor(
----------------
joker-eph wrote:

A shortcut condition before this one can be `reshape->getParentBlock() == dim->getParentBlock()`

https://github.com/llvm/llvm-project/pull/84225