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

[llvmlistbot at llvm.org]

Author: Federico Bruzzone
Date: 2026-09-17T14:44:49+02:00
New Revision: b5885865c07a5e6a20dac79318c23a0725ee90e1

URL: https://github.com/llvm/llvm-project/commit/b5885865c07a5e6a20dac79318c23a0725ee90e1
DIFF: https://github.com/llvm/llvm-project/commit/b5885865c07a5e6a20dac79318c23a0725ee90e1.diff

LOG: [mlir][bufferization] Fix 32-bit overflow in `PromoteBuffersToStack` threshold check (#223986)

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

Fixes #223937

---------

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
    mlir/test/Transforms/promote-buffers-to-stack.mlir

Removed: 
    


################################################################################
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.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
 
 // -----