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

[Federico Bruzzone]

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

This PR fix 32-bit overflow in `PromoteBuffersToStack` threshold check.

Fixes #223937

>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] [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
+}