[mlir][bufferization] Fix promote-buffers-to-stack crash caused by nested memref

[CoTinker]

This PR fixes a crash caused by nested memref type. Fixes #61375.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-bufferization

Author: Longsheng Mou (CoTinker)

Changes

This PR fixes a crash cause by nested memref type. Fixes #61375.

---

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

2 Files Affected:

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

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<ShapedType>(alloc.getType());
   if (!type || !alloc.getDefiningOp<memref::AllocOp>())
     return false;
+  auto elementType = type.getElementType();
+  // Not support nested memref type.
+  if (isa<MemRefType>(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<index, 256>>} {
 // 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
+}

[CoTinker]

Ping~

[amrami]

LGTM
I have no permissions to approve

[matthias-springer]

This will still crash on UnrankedMemRefType and certain types that implement MemRefElementTypeInterface. I think the right way to fix this is to add a helper function DataLayout::hasKnownTypeSize(Type). Internally, there could be function that returns std::optional<TypeSize>, which is called by both getTypeSizeInBits and hasKnownTypeSize.

[CoTinker]

Great idea.