[Mlir-commits] [mlir] [mlir][bufferization] Fix promote-buffers-to-stack crash cause by nested memref (PR #127441)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-bufferization

Author: Longsheng Mou (CoTinker)

<details>
<summary>Changes</summary>

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
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
+}

``````````

</details>


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