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

[Matthias Springer]

================
@@ -1069,39 +1069,6 @@ 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 failure();
-
-    // Place the load directly after the reshape to ensure that the shape memref
-    // was not mutated.
-    rewriter.setInsertionPointAfter(reshape);
-    Location loc = dim.getLoc();
-    Value load =
-        rewriter.create<LoadOp>(loc, reshape.getShape(), dim.getIndex());
----------------
matthias-springer wrote:

To avoid the dominance error: Instead of full `DominanceInfo` you could do a cheaper check here that:
* `dim.getIndex()` is defined in a parent block of `reshape`.
* Or: `dim.getIndex()` is defined in the same block as `reshape` but before `reshape`.

The annoying part is that `dim.getIndex()` could be a block argument, so there's one more special case. I've have into use cases where I needed this kind of check (defined in parent op or in the same block before a given op) many times, so we could add a such a helper function to `Dominance.h`.

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