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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-memref

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/183995.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+6-2) 
- (modified) mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir (+19) 


``````````diff
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/Vector/vector-multi-reduction-pass-lowering.mlir b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
index ddbc5c7bdb2c0..262bf526a1769 100644
--- a/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
@@ -61,3 +61,22 @@ func.func @vector_multi_reduction_masked(%arg0: vector<2x4xf32>, %acc: vector<2x
 //  INNER-PARALLEL: %[[INNERMASK:.+]] = vector.extract %[[TPMASK]][0] : vector<2xi1> from vector<4x2xi1>
 //  INNER-PARALLEL: %[[REDUCED:.+]] = arith.mulf %[[INNERVEC]], %[[ACC]] : vector<2xf32>
 //  INNER-PARALLEL: arith.select %[[INNERMASK]], %[[REDUCED]], %[[ACC]] : vector<2xi1>, vector<2xf32>
+
+// -----
+
+// Verify that -lower-vector-multi-reduction does not crash when a memref.dim
+// on a subview has 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
+
+// ALL-LABEL: func @no_crash_dim_of_ambiguous_subview
+// ALL: 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: ?>>
+  vector.print %dim : index
+  return
+}

``````````

</details>


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