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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/186426

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

>From 66c2cc7043c179050c261cf550248781bf46dd58 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 13 Mar 2026 07:04:26 -0700
Subject: [PATCH] [mlir][bufferization] Fix crash in promote-buffers-to-stack
 for nested memrefs

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
---
 .../Transforms/BufferOptimizations.cpp             |  9 ++++++++-
 mlir/test/Transforms/promote-buffers-to-stack.mlir | 14 ++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

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