[Mlir-commits] [mlir] [mlir][bufferization] Fix integer overflow crash in promote-buffers-to-stack (PR #186276)

[Matthias Springer]

================
@@ -105,7 +105,17 @@ static bool defaultIsSmallAlloc(Value alloc, unsigned maximumSizeInBytes,
   }
   unsigned bitwidth = mlir::DataLayout::closest(alloc.getDefiningOp())
                           .getTypeSizeInBits(type.getElementType());
-  return type.getNumElements() * bitwidth <= maximumSizeInBytes * 8;
+  // Use tryGetNumElements to avoid an assertion on integer overflow (e.g. for
+  // very large statically-shaped memrefs).  If the element count overflows
+  // int64_t the allocation is certainly not "small", so return false.
+  std::optional<int64_t> numElements = type.tryGetNumElements();
+  if (!numElements)
+    return false;
+  // Guard against overflow in the size computation as well.
+  if (bitwidth != 0 &&
+      *numElements > (int64_t)(maximumSizeInBytes * 8ULL / bitwidth))
----------------
matthias-springer wrote:

nit: C++-style cast

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