diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp index 93c1f9a4f2b55..5420288fe0722 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp @@ -89,6 +89,11 @@ static bool defaultIsSmallAlloc(Value alloc, unsigned maximumSizeInBytes, auto type = dyn_cast(alloc.getType()); if (!type || !alloc.getDefiningOp()) return false; + auto elementType = type.getElementType(); + // Not support nested memref type. + if (isa(elementType)) + return false; + if (!type.hasStaticShape()) { // Check if the dynamic shape dimension of the alloc is produced by // `memref.rank`. If this is the case, it is likely to be small. @@ -104,7 +109,7 @@ static bool defaultIsSmallAlloc(Value alloc, unsigned maximumSizeInBytes, return false; } unsigned bitwidth = mlir::DataLayout::closest(alloc.getDefiningOp()) - .getTypeSizeInBits(type.getElementType()); + .getTypeSizeInBits(elementType); return type.getNumElements() * bitwidth <= maximumSizeInBytes * 8; } diff --git a/mlir/test/Transforms/promote-buffers-to-stack.mlir b/mlir/test/Transforms/promote-buffers-to-stack.mlir index f7f2d2ec114ca..2648264458e59 100644 --- a/mlir/test/Transforms/promote-buffers-to-stack.mlir +++ b/mlir/test/Transforms/promote-buffers-to-stack.mlir @@ -611,3 +611,14 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry>} { // LOWLIMIT-NEXT: memref.alloc() {alignment = 64 : i64, custom_attr} // RANK-NEXT: memref.alloca() {alignment = 64 : i64, custom_attr} // CHECK-NEXT: return + +// ----- + +// Test Case: AllocOp with nested memref type, it is not +// converted. Ensure this case not crash. + +// CHECK-LABEL: func @nested_memref +func.func @nested_memref() { + %0 = memref.alloc() : memref<1xmemref<1xf32>> + return +}