[Mlir-commits] [mlir] [mlir][bufferization] Fix 32-bit overflow in `PromoteBuffersToStack` threshold check (PR #223986)

[Federico Bruzzone]

https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/223986

>From 563a20a44538b5c5bf850aa298df0b245d80d8d8 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 16 Sep 2026 11:40:54 +0200
Subject: [PATCH 1/2] [mlir][bufferization] Fix 32-bit overflow in
 PromoteBuffersToStack threshold check

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../Bufferization/Transforms/BufferOptimizations.cpp  |  3 ++-
 .../promote-buffers-to-stack-threshold-overflow.mlir  | 11 +++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Transforms/promote-buffers-to-stack-threshold-overflow.mlir

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 21b9ed446fb29..4f2a4734056b6 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -122,7 +122,8 @@ static bool defaultIsSmallAlloc(Value alloc, unsigned maximumSizeInBytes,
   if (bitwidth != 0 &&
       *numElements > static_cast<int64_t>(maximumSizeInBytes * 8ULL / bitwidth))
     return false;
-  return *numElements * bitwidth <= maximumSizeInBytes * 8;
+  return *numElements * bitwidth <=
+         static_cast<int64_t>(maximumSizeInBytes * 8ULL);
 }
 
 /// Checks whether the given aliases leave the allocation scope.
diff --git a/mlir/test/Transforms/promote-buffers-to-stack-threshold-overflow.mlir b/mlir/test/Transforms/promote-buffers-to-stack-threshold-overflow.mlir
new file mode 100644
index 0000000000000..b0da7b971a709
--- /dev/null
+++ b/mlir/test/Transforms/promote-buffers-to-stack-threshold-overflow.mlir
@@ -0,0 +1,11 @@
+// RUN: mlir-opt -promote-buffers-to-stack="max-alloc-size-in-bytes=536870912" %s | FileCheck %s
+
+// CHECK-LABEL: func @thresholdOverflow
+func.func @thresholdOverflow(%value: i8) -> i8 {
+  %c0 = arith.constant 0 : index
+  // CHECK: memref.alloca()
+  %buffer = memref.alloc() : memref<1xi8>
+  memref.store %value, %buffer[%c0] : memref<1xi8>
+  %result = memref.load %buffer[%c0] : memref<1xi8>
+  return %result : i8
+}

>From 2db65387b713dcc32dd205063edfb1c1e9134180 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Thu, 17 Sep 2026 07:46:56 +0200
Subject: [PATCH 2/2] Address comment

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../promote-buffers-to-stack-threshold-overflow.mlir | 11 -----------
 mlir/test/Transforms/promote-buffers-to-stack.mlir   | 12 ++++++++++--
 2 files changed, 10 insertions(+), 13 deletions(-)
 delete mode 100644 mlir/test/Transforms/promote-buffers-to-stack-threshold-overflow.mlir

diff --git a/mlir/test/Transforms/promote-buffers-to-stack-threshold-overflow.mlir b/mlir/test/Transforms/promote-buffers-to-stack-threshold-overflow.mlir
deleted file mode 100644
index b0da7b971a709..0000000000000
--- a/mlir/test/Transforms/promote-buffers-to-stack-threshold-overflow.mlir
+++ /dev/null
@@ -1,11 +0,0 @@
-// RUN: mlir-opt -promote-buffers-to-stack="max-alloc-size-in-bytes=536870912" %s | FileCheck %s
-
-// CHECK-LABEL: func @thresholdOverflow
-func.func @thresholdOverflow(%value: i8) -> i8 {
-  %c0 = arith.constant 0 : index
-  // CHECK: memref.alloca()
-  %buffer = memref.alloc() : memref<1xi8>
-  memref.store %value, %buffer[%c0] : memref<1xi8>
-  %result = memref.load %buffer[%c0] : memref<1xi8>
-  return %result : i8
-}
diff --git a/mlir/test/Transforms/promote-buffers-to-stack.mlir b/mlir/test/Transforms/promote-buffers-to-stack.mlir
index 3e4652b6acd52..4c8a5659e5866 100644
--- a/mlir/test/Transforms/promote-buffers-to-stack.mlir
+++ b/mlir/test/Transforms/promote-buffers-to-stack.mlir
@@ -1,6 +1,7 @@
 // RUN: mlir-opt -promote-buffers-to-stack -split-input-file %s | FileCheck %s --check-prefix=CHECK --check-prefix DEFINDEX
 // RUN: mlir-opt -promote-buffers-to-stack="max-alloc-size-in-bytes=64" -split-input-file %s | FileCheck %s --check-prefix=CHECK --check-prefix LOWLIMIT
 // RUN: mlir-opt -promote-buffers-to-stack="max-rank-of-allocated-memref=2" -split-input-file %s | FileCheck %s --check-prefix=CHECK --check-prefix RANK
+// RUN: mlir-opt -promote-buffers-to-stack="max-alloc-size-in-bytes=536870912" -split-input-file %s | FileCheck %s --check-prefix=CHECK --check-prefix HIGHLIMIT
 
 // This file checks the behavior of PromoteBuffersToStack pass for converting
 // AllocOps into AllocaOps, if possible.
@@ -570,7 +571,9 @@ func.func @loop_nested_if_alloc(
 // -----
 
 // Test Case: The allocated buffer is too large and, hence, it is not
-// converted. In the actual implementation the largest size is 1KB.
+// converted. In the actual implementation the largest size is 1KB. Under a
+// high limit (e.g. 512 MiB) that overflows `maxAllocSizeInBytes * 8` as a
+// 32-bit unsigned value, it must still be converted.
 
 // CHECK-LABEL: func @large_buffer_allocation
 func.func @large_buffer_allocation(%arg0: memref<2048xf32>) {
@@ -579,7 +582,10 @@ func.func @large_buffer_allocation(%arg0: memref<2048xf32>) {
   return
 }
 
-// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// DEFINDEX-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// LOWLIMIT-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// RANK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// HIGHLIMIT-NEXT: %[[ALLOC:.*]] = memref.alloca()
 // CHECK-NEXT: test.copy
 
 // -----
@@ -596,6 +602,7 @@ func.func @indexElementType() {
 // DEFINDEX-NEXT: memref.alloca()
 // LOWLIMIT-NEXT: memref.alloca()
 // RANK-NEXT: memref.alloca()
+// HIGHLIMIT-NEXT: memref.alloca()
 // CHECK-NEXT: return
 
 // -----
@@ -610,6 +617,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<index, 256>>} {
 // DEFINDEX-NEXT: memref.alloca() alignment = 64 {custom_attr}
 // LOWLIMIT-NEXT: memref.alloc() alignment = 64 {custom_attr}
 // RANK-NEXT: memref.alloca() alignment = 64 {custom_attr}
+// HIGHLIMIT-NEXT: memref.alloca() alignment = 64 {custom_attr}
 // CHECK-NEXT: return
 
 // -----