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

[Longsheng Mou]

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

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

>From 403cc1ee1b54449bd10f03bf107c01c62820344e Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Mon, 17 Feb 2025 11:51:36 +0800
Subject: [PATCH 1/2] [mlir][bufferization] Fix `promote-buffers-to-stack`
 crash cause by nested memref

This PR fix a crash cause by nested memref type.
---
 .../Bufferization/Transforms/BufferOptimizations.cpp       | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

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

>From 0df1d0f9416036861589fb07634c8e45a8a5f6fb Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Mon, 17 Feb 2025 14:11:19 +0800
Subject: [PATCH 2/2] Update promote-buffers-to-stack.mlir

---
 mlir/test/Transforms/promote-buffers-to-stack.mlir | 11 +++++++++++
 1 file changed, 11 insertions(+)

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