[Mlir-commits] [mlir] 5c3d0dd - [mlir][memref] Fix crash in DimOp::fold for subview with ambiguous dropped dims (#183995)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-05T17:04:58+01:00
New Revision: 5c3d0dd7a482b38e5a68f462d7ea874f55bc67f6

URL: https://github.com/llvm/llvm-project/commit/5c3d0dd7a482b38e5a68f462d7ea874f55bc67f6
DIFF: https://github.com/llvm/llvm-project/commit/5c3d0dd7a482b38e5a68f462d7ea874f55bc67f6.diff

LOG: [mlir][memref] Fix crash in DimOp::fold for subview with ambiguous dropped dims (#183995)

When a rank-reducing subview has multiple size-1 source dimensions and
all strides are dynamic, `computeMemRefRankReductionMask` cannot
unambiguously determine which dimension is dropped. It would fall back
to an arbitrary choice that could be wrong.

In `DimOp::fold`, the code then maps a result dimension to that
(possibly incorrect) source dimension and asserts
`isDynamicSize(sourceIndex)`. If `getDroppedDims()` picked the wrong
dropped dim, the source dim at `sourceIndex` may be statically sized,
causing the assertion to fire.

Fix this by replacing the assert with a guard that returns `{}` (no
fold) when the resolved source dimension is not dynamic. This is safe
because the fold already handles the statically-sized result case
earlier.

Fixes #111244

Assisted-by: Claude Code

Added: 
    

Modified: 
    mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
    mlir/test/Dialect/MemRef/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index d36b72d5652c9..404b2aacf1450 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -1150,22 +1150,20 @@ OpFoldResult DimOp::fold(FoldAdaptor adaptor) {
              memrefType.getDynamicDimIndex(unsignedIndex));
 
   if (auto subview = dyn_cast_or_null<SubViewOp>(definingOp)) {
-    llvm::SmallBitVector unusedDims = subview.getDroppedDims();
-    unsigned resultIndex = 0;
-    unsigned sourceRank = subview.getSourceType().getRank();
-    unsigned sourceIndex = 0;
-    for (auto i : llvm::seq<unsigned>(0, sourceRank)) {
-      if (unusedDims.test(i))
+    // The result dim is dynamic (the static case was handled above). Dropped
+    // dims always have static size 1, so dynamic source sizes are never
+    // dropped and map in order to the dynamic result dims. Find the k-th
+    // dynamic source size, where k is the dynamic dim index of the result dim.
+    unsigned dynamicResultDimIdx = memrefType.getDynamicDimIndex(unsignedIndex);
+    unsigned dynamicIdx = 0;
+    for (OpFoldResult size : subview.getMixedSizes()) {
+      if (llvm::isa<Attribute>(size))
         continue;
-      if (resultIndex == unsignedIndex) {
-        sourceIndex = i;
-        break;
-      }
-      resultIndex++;
+      if (dynamicIdx == dynamicResultDimIdx)
+        return size;
+      dynamicIdx++;
     }
-    assert(subview.isDynamicSize(sourceIndex) &&
-           "expected dynamic subview size");
-    return subview.getDynamicSize(sourceIndex);
+    return {};
   }
 
   // dim(memrefcast) -> dim

diff  --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir
index f0193d066dc37..92754dc919695 100644
--- a/mlir/test/Dialect/MemRef/canonicalize.mlir
+++ b/mlir/test/Dialect/MemRef/canonicalize.mlir
@@ -1639,3 +1639,24 @@ func.func @non_replace_view_negative_static_dims(%src: memref<?xi8>, %offset : i
   %res = memref.view %src[%offset][%c-1] : memref<?xi8> to memref<?x4xi32>
   return %res : memref<?x4xi32>
 }
+
+// -----
+
+// Verify that canonicalization does not crash when a memref.dim is applied to
+// a subview with ambiguous dropped dimensions (multiple size-1 source dims with
+// all-dynamic strides). The dim should be folded to the corresponding subview
+// size operand.
+// See: https://github.com/llvm/llvm-project/issues/111244
+
+// CHECK-LABEL: func @no_crash_dim_of_ambiguous_subview
+// CHECK-SAME: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: index) -> index
+// CHECK-NOT: memref.dim
+// CHECK: return %[[ARG1]]
+func.func @no_crash_dim_of_ambiguous_subview(
+    %arg0: memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>, %arg1: index) -> index {
+  %c1 = arith.constant 1 : index
+  %subview = memref.subview %arg0[0, 0, 0] [1, %arg1, 1] [1, 1, 1]
+      : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> to memref<1x?xf32, strided<[?, ?], offset: ?>>
+  %dim = memref.dim %subview, %c1 : memref<1x?xf32, strided<[?, ?], offset: ?>>
+  return %dim : index
+}