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

[Sayan Saha]

================
@@ -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()) {
----------------
sahas3 wrote:

I was hoping that the messages in the `notifyMatchFailure` will inform the intention of the two conditions. I've added some more comments to make the intention clearer.

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