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

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/183995

>From 9883a70fe8bba738da97d776b8c28d0e9a122dea Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sun, 1 Mar 2026 02:18:38 -0800
Subject: [PATCH] [mlir][memref] Fix crash in DimOp::fold for subview with
 ambiguous dropped dims

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
---
 mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp   |  8 ++++++--
 mlir/test/Dialect/MemRef/canonicalize.mlir | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 844e6183cff06..52a2015793cac 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -1084,8 +1084,12 @@ OpFoldResult DimOp::fold(FoldAdaptor adaptor) {
       }
       resultIndex++;
     }
-    assert(subview.isDynamicSize(sourceIndex) &&
-           "expected dynamic subview size");
+    // If the source dimension is not dynamic, getDroppedDims() may have
+    // computed an incorrect mapping (e.g., when strides are all dynamic and
+    // there are multiple unit-size candidate dimensions), so conservatively
+    // bail out.
+    if (!subview.isDynamicSize(sourceIndex))
+      return {};
     return subview.getDynamicSize(sourceIndex);
   }
 
diff --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir
index 3cfea1e8cd961..8476e226968a5 100644
--- a/mlir/test/Dialect/MemRef/canonicalize.mlir
+++ b/mlir/test/Dialect/MemRef/canonicalize.mlir
@@ -1639,3 +1639,21 @@ 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 left unfolded.
+// See: https://github.com/llvm/llvm-project/issues/111244
+
+// CHECK-LABEL: func @no_crash_dim_of_ambiguous_subview
+// CHECK: memref.dim
+func.func @no_crash_dim_of_ambiguous_subview(
+    %arg0: memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>, %arg1: 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
+}