[Mlir-commits] [mlir] [mlir][memref] Avoid overflow in mem2reg alloca checks (PR #205245)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Aug 13 20:00:29 PDT 2026
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/6] [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/6] [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/6] 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;
>From 98fa7cf8709b6e08d335c51edc5d3ce39c545cb8 Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Wed, 24 Jun 2026 10:08:46 +0800
Subject: [PATCH 4/6] Add regression test for mem2reg overflow fix in
MemRefMemorySlot.cpp (bug #204297 / #205245)
---
.../test/Dialect/MemRef/mem2reg-overflow.mlir | 22 +++++++++++++++++++
1 file changed, 22 insertions(+)
create mode 100644 mlir/test/Dialect/MemRef/mem2reg-overflow.mlir
diff --git a/mlir/test/Dialect/MemRef/mem2reg-overflow.mlir b/mlir/test/Dialect/MemRef/mem2reg-overflow.mlir
new file mode 100644
index 0000000000000..4b0f224bdbe98
--- /dev/null
+++ b/mlir/test/Dialect/MemRef/mem2reg-overflow.mlir
@@ -0,0 +1,22 @@
+# RUN: mlir-opt %s -mem2reg -verify-diagnostics | FileCheck %s
+
+# CHECK: module
+module {
+ // This used to crash mlir-opt with the assertion in BuiltinTypeInterfaces.cpp:86
+ // (9223372036854775807 * 3 overflowing)
+ func.func @alloca_unconvertable_memory_space() {
+ return
+ ^bb1(%0: index):
+ %alloca = memref.alloca() : memref<9223372036854775807x3xi32>
+ return
+ }
+
+ // This is a copy of the crashing input from the bug report
+ // It should now run successfully with no assertion and no crash
+ func.func @test_mem2reg_works() {
+ %0 = arith.constant 42 : i32
+ %1 = memref.alloca() : memref<9223372036854775807x3xi32>
+ memref.store %0, %1[%0] : memref<9223372036854775807x3xi32>
+ return
+ }
+}
>From ecada4ce75de9fb54adad4a97601ea7d4013ab9a Mon Sep 17 00:00:00 2001
From: LouisLu060211 <38174270+LouisLu060211 at users.noreply.github.com>
Date: Fri, 14 Aug 2026 10:53:04 +0800
Subject: [PATCH 5/6] Delete mlir/test/Dialect/MemRef/mem2reg-overflow.mlir
---
.../test/Dialect/MemRef/mem2reg-overflow.mlir | 22 -------------------
1 file changed, 22 deletions(-)
delete mode 100644 mlir/test/Dialect/MemRef/mem2reg-overflow.mlir
diff --git a/mlir/test/Dialect/MemRef/mem2reg-overflow.mlir b/mlir/test/Dialect/MemRef/mem2reg-overflow.mlir
deleted file mode 100644
index 4b0f224bdbe98..0000000000000
--- a/mlir/test/Dialect/MemRef/mem2reg-overflow.mlir
+++ /dev/null
@@ -1,22 +0,0 @@
-# RUN: mlir-opt %s -mem2reg -verify-diagnostics | FileCheck %s
-
-# CHECK: module
-module {
- // This used to crash mlir-opt with the assertion in BuiltinTypeInterfaces.cpp:86
- // (9223372036854775807 * 3 overflowing)
- func.func @alloca_unconvertable_memory_space() {
- return
- ^bb1(%0: index):
- %alloca = memref.alloca() : memref<9223372036854775807x3xi32>
- return
- }
-
- // This is a copy of the crashing input from the bug report
- // It should now run successfully with no assertion and no crash
- func.func @test_mem2reg_works() {
- %0 = arith.constant 42 : i32
- %1 = memref.alloca() : memref<9223372036854775807x3xi32>
- memref.store %0, %1[%0] : memref<9223372036854775807x3xi32>
- return
- }
-}
>From e95126daf2e8c682c7c0031c95999650a35b28a8 Mon Sep 17 00:00:00 2001
From: LouisLu060211 <38174270+LouisLu060211 at users.noreply.github.com>
Date: Fri, 14 Aug 2026 11:00:15 +0800
Subject: [PATCH 6/6] Add test for alloca element count overflow
Add test for mem2reg with alloca element count overflow
---
mlir/test/Dialect/MemRef/mem2reg.mlir | 1 +
1 file changed, 1 insertion(+)
diff --git a/mlir/test/Dialect/MemRef/mem2reg.mlir b/mlir/test/Dialect/MemRef/mem2reg.mlir
index 966675d0941f2..db1d02f14b8a4 100644
--- a/mlir/test/Dialect/MemRef/mem2reg.mlir
+++ b/mlir/test/Dialect/MemRef/mem2reg.mlir
@@ -314,6 +314,7 @@ func.func @two_consecutive_merge_points(%cond1: i1, %cond2: i1) -> 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
More information about the Mlir-commits
mailing list