[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 2a0b4acec108e058902520dc00dac5b95ba6c60c Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Tue, 23 Jun 2026 11:18:32 +0800
Subject: [PATCH 1/3] [mlir][memref] Avoid overflow in mem2reg alloca checks

---
 mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index 6748e2cf71804..69cffde01470e 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -67,7 +67,8 @@ 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 +291,10 @@ 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;

>From f6d09dea23b5c0b62029ed07a350a9fae14dd664 Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Tue, 23 Jun 2026 11:18:32 +0800
Subject: [PATCH 2/3] [mlir][memref] Avoid overflow in mem2reg alloca checks

---
 mlir/test/Dialect/MemRef/mem2reg.mlir | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/mlir/test/Dialect/MemRef/mem2reg.mlir b/mlir/test/Dialect/MemRef/mem2reg.mlir
index 8f937c4efe75e..966675d0941f2 100644
--- a/mlir/test/Dialect/MemRef/mem2reg.mlir
+++ b/mlir/test/Dialect/MemRef/mem2reg.mlir
@@ -309,3 +309,16 @@ 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
+}

>From 743c69eba577f782aaa103b46a548873b881591a Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Wed, 24 Jun 2026 09:57:45 +0800
Subject: [PATCH 3/3] Your commit message here

---
 mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index 69cffde01470e..aa95ebf540fd7 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -67,7 +67,8 @@ SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
   if (!type.hasStaticShape())
     return {};
   // Make sure the memref contains only a single element.
-  std::optional<int64_t> numElements = ShapedType::tryGetNumElements(type.getShape());
+  std::optional<int64_t> numElements = 
+      ShapedType::tryGetNumElements(type.getShape());
   if (!numElements || *numElements != 1)
     return {};
 
@@ -293,8 +294,10 @@ struct MemRefDestructurableTypeExternalModel
     constexpr int64_t maxMemrefSizeForDestructuring = 16;
     if (!memrefType.hasStaticShape())
       return {};
-    std::optional<int64_t> numElements = ShapedType::tryGetNumElements(memrefType.getShape());
-    if (!numElements || *numElements > maxMemrefSizeForDestructuring || *numElements == 1)
+    std::optional<int64_t> numElements = 
+        ShapedType::tryGetNumElements(memrefType.getShape());
+    if (!numElements || *numElements > maxMemrefSizeForDestructuring || 
+        *numElements == 1)
       return {};
 
     DenseMap<Attribute, Type> destructured;