[Mlir-commits] [mlir] 200e8c5 - [MLIR][MemRef] Fix DimOfReifyRankedShapedTypeOpInterface IR-change on failure (#188973)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-04-13T12:20:29+02:00
New Revision: 200e8c589e41f561b8dd7ed7c9b5e5ba26ab8300

URL: https://github.com/llvm/llvm-project/commit/200e8c589e41f561b8dd7ed7c9b5e5ba26ab8300
DIFF: https://github.com/llvm/llvm-project/commit/200e8c589e41f561b8dd7ed7c9b5e5ba26ab8300.diff

LOG: [MLIR][MemRef] Fix DimOfReifyRankedShapedTypeOpInterface IR-change on failure (#188973)

DimOfReifyRankedShapedTypeOpInterface::matchAndRewrite called
reifyDimOfResult via the PatternRewriter. Some implementations delegate
to the coarse-grained reifyResultShapes, which creates ops for ALL
dimensions (e.g. a tensor.dim) before discovering that a specific
dimension is not reifiable (signalled by an empty OpFoldResult).

The pattern then returned failure() once it saw the empty OpFoldResult,
but the newly created ops were already in the IR. Under
MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS this triggered "pattern
returned failure but IR did change".

Fix: record the op immediately before the matched dim op, so we can
identify ops inserted during the reification attempt. If reification
returns an empty (unreifiable) OpFoldResult, erase those newly created
ops before returning failure, restoring the IR to its original state.

Assisted-by: Claude Code

Added: 
    

Modified: 
    mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp
    mlir/test/Interfaces/InferShapedTypeOpInterface/resolve-shaped-type-result-dims.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp b/mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp
index c498c8a60bf6e..90a58cc1a807f 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp
@@ -90,13 +90,39 @@ struct DimOfReifyRankedShapedTypeOpInterface : public OpRewritePattern<OpTy> {
     if (!dimIndex)
       return failure();
 
+    // Save the op immediately before dimOp so we can identify and erase any
+    // ops inserted during the reification attempt if it fails. The
+    // pattern-rewrite invariant requires the IR to be unchanged on failure.
+    Operation *opBeforeReify = dimOp->getPrevNode();
+
+    // Erase any ops inserted between opBeforeReify and dimOp in reverse order
+    // to respect use-def chains within that range. Collect pointers first to
+    // avoid iterator invalidation: erasing a node in an ilist invalidates
+    // iterators to that node, and std::reverse_iterator stores the iterator to
+    // the *next* forward element, so make_early_inc_range(reverse(...)) would
+    // still dereference a stale iterator after erasure.
+    auto eraseInsertedOps = [&]() {
+      Block::iterator begin = opBeforeReify
+                                  ? std::next(opBeforeReify->getIterator())
+                                  : dimOp->getBlock()->begin();
+      SmallVector<Operation *> toErase;
+      for (Block::iterator it = begin; it != dimOp->getIterator(); ++it)
+        toErase.push_back(&*it);
+      for (Operation *op : llvm::reverse(toErase))
+        rewriter.eraseOp(op);
+    };
+
     FailureOr<OpFoldResult> replacement = reifyDimOfResult(
         rewriter, dimValue.getOwner(), dimValue.getResultNumber(), *dimIndex);
-    if (failed(replacement))
-      return failure();
-    // Check if the OpFoldResult is empty (unreifiable dimension).
-    if (!replacement.value())
+    // An empty (or failed) OpFoldResult signals that this specific dimension
+    // cannot be reified. Some implementations materialize all dimensions at
+    // once (e.g. via reifyResultShapes) and may create ops for other dimensions
+    // before discovering that this dimension is not reifiable. Erase those
+    // stray ops before returning failure.
+    if (failed(replacement) || !replacement.value()) {
+      eraseInsertedOps();
       return failure();
+    }
     Value replacementVal = getValueOrCreateConstantIndexOp(
         rewriter, dimOp.getLoc(), replacement.value());
     rewriter.replaceOp(dimOp, replacementVal);

diff  --git a/mlir/test/Interfaces/InferShapedTypeOpInterface/resolve-shaped-type-result-dims.mlir b/mlir/test/Interfaces/InferShapedTypeOpInterface/resolve-shaped-type-result-dims.mlir
index 624e0990a4bb3..f41312839c094 100644
--- a/mlir/test/Interfaces/InferShapedTypeOpInterface/resolve-shaped-type-result-dims.mlir
+++ b/mlir/test/Interfaces/InferShapedTypeOpInterface/resolve-shaped-type-result-dims.mlir
@@ -178,3 +178,24 @@ func.func @test_unreifiable_dim_of_result_shape(%arg0 : tensor<?x?xf32>)
 //   CHECK-DAG:   %[[OP:.+]] = "test.unreifiable_dim_of_result_shape"(%[[ARG0]])
 //       CHECK:   %[[D1:.+]] = tensor.dim %[[OP]], %[[C1]]
 //       CHECK:   return %[[D0]], %[[D1]]
+
+// -----
+
+// Regression test: verify that when reifyResultShapes creates ops for dim 0
+// but signals dim 1 is not reifiable (empty OpFoldResult), those stray ops are
+// erased before failure is returned. Without the fix, the stray tensor.dim op
+// on %arg0 would remain in the IR (caught by MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS).
+func.func @test_unreifiable_result_shapes_no_stray_ops(%arg0 : tensor<?x?xf32>)
+    -> index {
+  %c1 = arith.constant 1 : index
+  %0 = "test.unreifiable_result_shapes"(%arg0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
+  %d1 = tensor.dim %0, %c1 : tensor<?x?xf32>
+  return %d1 : index
+}
+// CHECK-LABEL: func @test_unreifiable_result_shapes_no_stray_ops(
+//  CHECK-SAME:   %[[ARG0:.+]]: tensor<?x?xf32>)
+//       CHECK:   %[[C1:.+]] = arith.constant 1 : index
+//       CHECK:   %[[OP:.+]] = "test.unreifiable_result_shapes"(%[[ARG0]])
+// CHECK-NOT:     tensor.dim %[[ARG0]]  // key: no stray dim on the input arg
+//       CHECK:   %[[D1:.+]] = tensor.dim %[[OP]], %[[C1]]
+//       CHECK:   return %[[D1]]