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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Federico Bruzzone (FedericoBruzzone)

<details>
<summary>Changes</summary>

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

Fixes #<!-- -->223937

---
Full diff: https://github.com/llvm/llvm-project/pull/223986.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp (+2-1) 
- (added) mlir/test/Transforms/promote-buffers-to-stack-threshold-overflow.mlir (+11) 


``````````diff
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
+}

``````````

</details>


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