[Mlir-commits] [mlir] [mlir][bufferization] Fix crash in buffer-hoisting on unreachable blocks (PR #184094)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-bufferization

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

`BufferAllocationHoisting::hoist()` iterates over all allocations in a function and tries to move them to their common dominator block. When an allocation is in an unreachable basic block (a block that cannot be reached from the function entry), the LLVM dominator tree has no node for that block. Two subsequent calls would then crash:

1. `findCommonDominator` calls `findNearestCommonDominator` which calls the LLVM DomTree equivalent and hits the assertion `NodeA && "A must be in the tree"`.

2. `findPlacementBlock` calls `dominators.getNode(currentBlock)` on the unreachable block, which returns null, causing a null dereference on `->getIDom()`.

Fix by skipping allocations whose parent block is not reachable from the entry block of the function. Such allocations are in dead code and there is no need to hoist them.

Fixes #<!-- -->118445

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp (+5) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir (+13) 


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index ebd0d827526d7..b14989b88351f 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -193,6 +193,11 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
         continue;
       Operation *definingOp = allocValue.getDefiningOp();
       assert(definingOp && "No defining op");
+      // Skip allocations in blocks that are not reachable from the function
+      // entry. Such blocks are dead code and the dominator tree analysis may
+      // not have nodes for them, which would cause crashes below.
+      if (!dominators.isReachableFromEntry(allocValue.getParentBlock()))
+        continue;
       auto operands = definingOp->getOperands();
       auto resultAliases = aliases.resolve(allocValue);
       // Determine the common dominator block of all aliases.
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir
index 33fc158f9f51a..4c07df951e280 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir
@@ -789,3 +789,16 @@ func.func @loop_nested_alloc_dyn_dependency(
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: {{.*}} = scf.for
 //      CHECK: %[[ALLOC1:.*]] = memref.alloc({{.*}})
+
+// Verify that --buffer-hoisting does not crash on allocations in unreachable
+// blocks (blocks with no predecessors that are not the function entry block).
+// See: https://github.com/llvm/llvm-project/issues/118445
+// CHECK-LABEL: func @unreachable_alloc
+func.func @unreachable_alloc(%arg0: f32) {
+  %c0 = arith.constant 0 : index
+  return
+^bb1(%0: i32):  // no predecessors
+  %alloc = memref.alloc() : memref<10xf32>
+  // CHECK: memref.alloc
+  return
+}

``````````

</details>


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