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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Sujian Hu (kezuo)

<details>
<summary>Changes</summary>

Fix #<!-- -->108364.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp (+7-2) 
- (added) mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-unreachable-block.mlir (+9) 


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 93c1f9a4f2b554..8e4491e4a3bdab 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 further up.
+        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
+}

``````````

</details>


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