[Mlir-commits] [mlir] [mlir][memref] Avoid overflow in mem2reg alloca checks (PR #205245)

[llvmlistbot at llvm.org]

https://github.com/LouisLu060211 updated https://github.com/llvm/llvm-project/pull/205245

>From f49fa875ee1f116cdbd35f9b6466576ccf8757fd Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Thu, 20 Aug 2026 09:53:22 +0800
Subject: [PATCH] [mlir][memref] Avoid overflow in mem2reg alloca checks

---
 .../Dialect/MemRef/IR/MemRefMemorySlot.cpp    | 20 ++++++++++++++-----
 mlir/test/Dialect/MemRef/mem2reg.mlir         | 12 +++++++++++
 mlir/test/Dialect/MemRef/sroa.mlir            | 12 +++++++++++
 3 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index cfcca90b2b969..0b1704b551fdd 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -96,8 +96,15 @@ SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
   MemRefType type = getType();
 
   // A single-element memref is promoted to a scalar SSA value.
-  if (type.hasStaticShape() && type.getNumElements() == 1)
-    return {MemorySlot{getResult(), type.getElementType()}};
+  if (type.hasStaticShape()) {
+    std::optional<int64_t> numElements =
+        ShapedType::tryGetNumElements(type.getShape());
+    // Element count overflow: not promotable.
+    if (!numElements)
+      return {};
+    if (*numElements == 1)
+      return {MemorySlot{getResult(), type.getElementType()}};
+  }
 
   // A multi-element memref can be promoted to a single vector SSA value when it
   // is only ever accessed as a whole buffer (e.g. through whole-buffer
@@ -345,9 +352,12 @@ struct MemRefDestructurableTypeExternalModel
   getSubelementIndexMap(Type type) const {
     auto memrefType = llvm::cast<MemRefType>(type);
     constexpr int64_t maxMemrefSizeForDestructuring = 16;
-    if (!memrefType.hasStaticShape() ||
-        memrefType.getNumElements() > maxMemrefSizeForDestructuring ||
-        memrefType.getNumElements() == 1)
+    if (!memrefType.hasStaticShape())
+      return {};
+    std::optional<int64_t> numElements =
+        ShapedType::tryGetNumElements(memrefType.getShape());
+    if (!numElements || *numElements > maxMemrefSizeForDestructuring ||
+        *numElements == 1)
       return {};
 
     DenseMap<Attribute, Type> destructured;
diff --git a/mlir/test/Dialect/MemRef/mem2reg.mlir b/mlir/test/Dialect/MemRef/mem2reg.mlir
index ba26b5f984a27..caf263fb1e05c 100644
--- a/mlir/test/Dialect/MemRef/mem2reg.mlir
+++ b/mlir/test/Dialect/MemRef/mem2reg.mlir
@@ -343,3 +343,15 @@ func.func @scalable_zero_extent_alloca() {
   %alloca = memref.alloca(%size) : memref<?xf32>
   return
 }
+
+// -----
+
+// Make sure mem2reg does not crash on an alloca whose element count overflows
+// int64. https://github.com/llvm/llvm-project/issues/204297
+
+// CHECK-LABEL: func.func @alloca_element_count_overflow
+func.func @alloca_element_count_overflow() {
+  // CHECK: memref.alloca() : memref<9223372036854775807x3xi32>
+  %alloca = memref.alloca() : memref<9223372036854775807x3xi32>
+  return
+}
diff --git a/mlir/test/Dialect/MemRef/sroa.mlir b/mlir/test/Dialect/MemRef/sroa.mlir
index 40ab9b3483b83..95b76ff5a0f78 100644
--- a/mlir/test/Dialect/MemRef/sroa.mlir
+++ b/mlir/test/Dialect/MemRef/sroa.mlir
@@ -173,3 +173,15 @@ func.func @no_out_of_bound_load(%arg0: i32, %arg1: i32) -> i32 {
   // CHECK: return %[[RES]] : i32
   return %res : i32
 }
+
+// -----
+
+// Make sure sroa does not crash on an alloca whose element count overflows
+// int64. Covers getSubelementIndexMap's tryGetNumElements check.
+
+// CHECK-LABEL: func.func @alloca_element_count_overflow
+func.func @alloca_element_count_overflow() {
+  // CHECK: memref.alloca() : memref<9223372036854775807x3xi32>
+  %alloca = memref.alloca() : memref<9223372036854775807x3xi32>
+  return
+}