[Mlir-commits] [mlir] [mlir][memref] Fix mem2reg crash on zero-extent alloca (PR #216851)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-memref

Author: Alessandro Potenza (alepot55)

<details>
<summary>Changes</summary>

`memref` allows an extent of `0`, but vector types require strictly positive
extents. Since the whole-buffer memref-to-vector promotion added in #<!-- -->211880,
`memref::AllocaOp::getPromotableSlots` builds a `VectorType` from the memref
shape, so a legal zero-extent alloca crashes `mlir-opt --mem2reg`:

```mlir
func.func @<!-- -->z() { %a = memref.alloca() : memref<0xf32>  return }
```

```
mlir-opt: .../StorageUniquerSupport.h:180: ... Assertion `succeeded(
  ConcreteT::verifyInvariants(getDefaultDiagnosticEmitFn(ctx), args...))' failed.
PLEASE submit a bug report ...
```

The alloca needs no uses at all: `getPromotableSlots` runs before the
`use_empty` filter in `tryToPromoteMemorySlots`.

Bail out of whole-buffer promotion when the shape contains a zero extent; such
a memref holds no elements, so there is nothing to promote. The scalable path
is guarded too: `matchVScaleMultiple` only tells us the extent is some constant
`C`, and `VectorType::verify` rejects `C <= 0`, so the condition is `*c > 0`
rather than `*c != 0`. The single-element scalar path is unaffected, since a
shape containing a zero extent can never have exactly one element.

`llvm::is_contained(getShape(), 0)` is deliberately used instead of
`getNumElements()`, so this does not overlap the element-count overflow fix in
flight in #<!-- -->205245 / #<!-- -->207670 (both of which are about `memref<9223372036854775807x3xi32>`
and leave this crash in place: `tryGetNumElements` returns `0` and falls through
to the same `VectorType::get`).

### Verified by execution, not by inspection

- The crash above reproduces on `d4e78d7f5`.
- The new tests fail on unpatched `mlir-opt` (abort, above) and pass with the patch.
- No regression: `mlir/test/Dialect/MemRef` 33/33, and `Dialect/SCF`, `Dialect/Affine`,
  `Dialect/Vector`, `Dialect/XeGPU`, `Transforms` all green (368 tests total).
- `git clang-format` reports no changes.

Caveat stated openly: the scalable path ships without a test, because exercising
it needs `vector.vscale` in a MemRef mem2reg test and I did not want to change
that file's pipeline. Happy to add one if a reviewer prefers.

---

Assisted-by: Claude (Anthropic)

This patch was written with AI assistance, disclosed per the LLVM AI Tool Use Policy.
Everything reported above as verified was verified by building and running, not by
asking a model whether it was true.


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


2 Files Affected:

- (modified) mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp (+8-2) 
- (modified) mlir/test/Dialect/MemRef/mem2reg.mlir (+22) 


``````````diff
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index a832b0c47809b..7a2bc0bbee0b8 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -103,15 +103,21 @@ SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
   // is only ever accessed as a whole buffer (e.g. through whole-buffer
   // `vector.transfer_read`/`vector.transfer_write`).
   if (VectorType::isValidElementType(type.getElementType())) {
+    // Vector types require strictly positive extents, so a memref with a zero
+    // extent (which holds no elements) has nothing to promote.
+    if (llvm::is_contained(type.getShape(), 0))
+      return {};
+
     // Static shape: a fixed-size vector of the same extents.
     if (type.hasStaticShape())
       return {MemorySlot{getResult(), VectorType::get(type.getShape(),
                                                       type.getElementType())}};
 
     // A 1-D memref whose single dynamic extent is `vector.vscale * C` maps to a
-    // scalable `vector<[C]x...>` slot.
+    // scalable `vector<[C]x...>` slot, for a strictly positive `C`.
     if (type.getRank() == 1 && type.isDynamicDim(0)) {
-      if (std::optional<int64_t> c = matchVScaleMultiple(getDynamicSizes()[0]))
+      if (std::optional<int64_t> c = matchVScaleMultiple(getDynamicSizes()[0]);
+          c && *c > 0)
         return {
             MemorySlot{getResult(), VectorType::get({*c}, type.getElementType(),
                                                     /*scalableDims=*/{true})}};
diff --git a/mlir/test/Dialect/MemRef/mem2reg.mlir b/mlir/test/Dialect/MemRef/mem2reg.mlir
index 8f937c4efe75e..add6bc2e457f0 100644
--- a/mlir/test/Dialect/MemRef/mem2reg.mlir
+++ b/mlir/test/Dialect/MemRef/mem2reg.mlir
@@ -309,3 +309,25 @@ func.func @two_consecutive_merge_points(%cond1: i1, %cond2: i1) -> i32 {
   // CHECK: return %[[RESULT]] : i32
   return %result : i32
 }
+
+// -----
+
+// A memref with a zero extent holds no elements and cannot be promoted: the
+// vector type mem2reg would build for whole-buffer promotion requires strictly
+// positive extents.
+
+// CHECK-LABEL: func.func @zero_extent_alloca
+func.func @zero_extent_alloca() {
+  // CHECK: memref.alloca() : memref<0xf32>
+  %alloca = memref.alloca() : memref<0xf32>
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @zero_extent_alloca_multi_dim
+func.func @zero_extent_alloca_multi_dim() {
+  // CHECK: memref.alloca() : memref<2x0x3xf32>
+  %alloca = memref.alloca() : memref<2x0x3xf32>
+  return
+}

``````````

</details>


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