[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 05ecb9ee9b71f8424a39a1f5012558c15f81122d 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

---
 mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp | 13 +++++++++----
 mlir/test/Dialect/MemRef/mem2reg.mlir           | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index 6748e2cf71804..aa95ebf540fd7 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -67,7 +67,9 @@ SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
   if (!type.hasStaticShape())
     return {};
   // Make sure the memref contains only a single element.
-  if (type.getNumElements() != 1)
+  std::optional<int64_t> numElements = 
+      ShapedType::tryGetNumElements(type.getShape());
+  if (!numElements || *numElements != 1)
     return {};
 
   return {MemorySlot{getResult(), type.getElementType()}};
@@ -290,9 +292,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 8f937c4efe75e..db1d02f14b8a4 100644
--- a/mlir/test/Dialect/MemRef/mem2reg.mlir
+++ b/mlir/test/Dialect/MemRef/mem2reg.mlir
@@ -309,3 +309,17 @@ func.func @two_consecutive_merge_points(%cond1: i1, %cond2: i1) -> i32 {
   // CHECK: return %[[RESULT]] : i32
   return %result : i32
 }
+
+// -----
+
+// 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() {
+  return
+^bb1(%0: index):
+  // CHECK: memref.alloca() : memref<9223372036854775807x3xi32>
+  %alloca = memref.alloca() : memref<9223372036854775807x3xi32>
+  return
+}