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

[Alessandro Potenza]

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

`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.


>From 3782460f038941fa618221cf84889f22c749747b Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Mon, 17 Aug 2026 22:56:39 +0200
Subject: [PATCH] [mlir][memref] Fix mem2reg crash on zero-extent alloca

`memref` allows an extent of 0, but vector types must have strictly
positive extents. When `memref::AllocaOp::getPromotableSlots` considers
whole-buffer promotion it builds a `VectorType` from the memref shape,
so a legal zero-extent alloca such as

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

crashes `mlir-opt --mem2reg` inside `VectorType::get`. 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.
Apply the same reasoning to the scalable path, where the extent comes
from a `vector.vscale * C` operand rather than the shape and `C` is only
known to be a constant: require it to be strictly positive.

The single-element scalar path is unaffected: a shape containing a zero
extent can never have exactly one element.
---
 .../Dialect/MemRef/IR/MemRefMemorySlot.cpp    | 10 +++++++--
 mlir/test/Dialect/MemRef/mem2reg.mlir         | 22 +++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

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
+}