[Mlir-commits] [mlir] [mlir][bufferization] Stop hoisting allocations beyond unreachable bl… (PR #113644)

[Sujian Hu]

https://github.com/kezuo updated https://github.com/llvm/llvm-project/pull/113644

>From b31f6a57d1e1eb75c3357e4b8c21b621cfc7570a Mon Sep 17 00:00:00 2001
From: hu shujian <zuoooke at gmail.com>
Date: Fri, 25 Oct 2024 11:04:35 +0800
Subject: [PATCH] [mlir][bufferization] Stop hoisting allocations beyond
 unreachable blocks

---
 .../Bufferization/Transforms/BufferOptimizations.cpp     | 9 +++++++--
 .../buffer-loop-hoisting-unreachable-block.mlir          | 9 +++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)
 create mode 100644 mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-unreachable-block.mlir

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 93c1f9a4f2b554..f733ddee95b558 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -250,8 +250,13 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
       DominanceInfoNode *idom = nullptr;
 
       // DominanceInfo doesn't support getNode queries for single-block regions.
-      if (!currentBlock->isEntryBlock())
-        idom = dominators.getNode(currentBlock)->getIDom();
+      if (!currentBlock->isEntryBlock()) {
+        auto domNode = dominators.getNode(currentBlock);
+        // If we encounter an unreachable block, we cannot move up further.
+        if (domNode == nullptr)
+          break;
+        idom = domNode->getIDom();
+      }
 
       if (idom && dominators.properlyDominates(parentBlock, idom->getBlock())) {
         // If the current immediate dominator is below the placement block, move
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-unreachable-block.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-unreachable-block.mlir
new file mode 100644
index 00000000000000..e396d3406c27fe
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-unreachable-block.mlir
@@ -0,0 +1,9 @@
+// RUN: mlir-opt -buffer-loop-hoisting %s
+
+func.func @unreachable_block() {
+  return
+^bb1:
+  %alloc = memref.alloc() : memref<2xf32>
+  memref.dealloc %alloc : memref<2xf32>
+  return
+}