[Mlir-commits] [mlir] [mlir][bufferization] Fix crash in promote-buffers-to-stack for nested memrefs (PR #186426)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-bufferization

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

The `--promote-buffers-to-stack` pass crashes when allocating a memref whose element type is itself a memref (e.g., `memref<1xmemref<2xf32>>`). This happens because `defaultIsSmallAlloc` calls
`DataLayout::getTypeSizeInBits` on the element type, but `MemRefType` (and other types without `DataLayoutTypeInterface`) trigger a fatal error when queried this way.

Fix the crash by checking whether the element type has data layout support before computing its size. Types that are not int/float, complex, index, or vector and do not implement `DataLayoutTypeInterface` are silently skipped (i.e., the allocation is not promoted).

Fixes #<!-- -->60092

Assisted-by: Claude Code

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp (+8-1) 
- (modified) mlir/test/Transforms/promote-buffers-to-stack.mlir (+14) 


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 0432a5799b440..21b9ed446fb29 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -19,6 +19,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Pass/Pass.h"
 
@@ -103,8 +104,14 @@ static bool defaultIsSmallAlloc(Value alloc, unsigned maximumSizeInBytes,
     }
     return false;
   }
+  // Skip if the element type has no data layout support (e.g., nested memrefs).
+  Type elemType = type.getElementType();
+  if (!elemType.isIntOrFloat() &&
+      !isa<ComplexType, IndexType, VectorType>(elemType) &&
+      !isa<DataLayoutTypeInterface>(elemType))
+    return false;
   unsigned bitwidth = mlir::DataLayout::closest(alloc.getDefiningOp())
-                          .getTypeSizeInBits(type.getElementType());
+                          .getTypeSizeInBits(elemType);
   // 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.
diff --git a/mlir/test/Transforms/promote-buffers-to-stack.mlir b/mlir/test/Transforms/promote-buffers-to-stack.mlir
index a71f7eaee1c04..c055236ce6dff 100644
--- a/mlir/test/Transforms/promote-buffers-to-stack.mlir
+++ b/mlir/test/Transforms/promote-buffers-to-stack.mlir
@@ -626,3 +626,17 @@ func.func @huge_static_memref() {
   %alloc = memref.alloc() : memref<3090540x3090540x3090540xi32>
   return
 }
+
+// -----
+
+// Test that allocations with memref element types (nested memrefs) are not
+// promoted to stack allocations, since no data layout information is available
+// for the inner memref type.
+
+// CHECK-LABEL: func @nestedMemref
+func.func @nestedMemref() {
+  %0 = memref.alloc() : memref<1xmemref<2xf32>>
+  return
+}
+// CHECK-NEXT: memref.alloc()
+// CHECK-NEXT: return

``````````

</details>


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