[mlir] Fix invalid hoisting of dependent allocs in buffer hoisting pass.

Buffer hoisting moves allocs upwards although it has dependency within its
nested region. This patch fixes this issue.

https://bugs.llvm.org/show_bug.cgi?id=49142

Differential Revision: https://reviews.llvm.org/D98248

diff --git a/mlir/lib/Transforms/BufferOptimizations.cpp b/mlir/lib/Transforms/BufferOptimizations.cpp
index 6c1e1fe..9ed4bb4 100644 (file)
--- a/mlir/lib/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Transforms/BufferOptimizations.cpp
@@ -142,13 +142,13 @@ public:
       // Check for additional allocation dependencies to compute an upper bound
       // for hoisting.
       Block *dependencyBlock = nullptr;
-      if (!operands.empty()) {
-        // If this node has dependencies, check all dependent nodes with respect
-        // to a common post dominator. This ensures that all dependency values
-        // have been computed before allocating the buffer.
-        ValueSetT dependencies(std::next(operands.begin()), operands.end());
-        dependencyBlock = findCommonDominator(*operands.begin(), dependencies,
-                                              postDominators);
+      // If this node has dependencies, check all dependent nodes. This ensures
+      // that all dependency values have been computed before allocating the
+      // buffer.
+      for (Value depValue : operands) {
+        Block *depBlock = depValue.getParentBlock();
+        if (!dependencyBlock || dominators.dominates(dependencyBlock, depBlock))
+          dependencyBlock = depBlock;
       }
 
       // Find the actual placement block and determine the start operation using
@@ -371,7 +371,7 @@ public:
   explicit PromoteBuffersToStackPass(std::function<bool(Value)> isSmallAlloc)
       : isSmallAlloc(std::move(isSmallAlloc)) {}
 
-  LogicalResult initialize(MLIRContext* context) override {
+  LogicalResult initialize(MLIRContext *context) override {
     if (isSmallAlloc == nullptr) {
       isSmallAlloc = [=](Value alloc) {
         return defaultIsSmallAlloc(alloc, maxAllocSizeInBytes,

diff --git a/mlir/test/Transforms/buffer-hoisting.mlir b/mlir/test/Transforms/buffer-hoisting.mlir
index 706c768..f45fffd 100644 (file)
--- a/mlir/test/Transforms/buffer-hoisting.mlir
+++ b/mlir/test/Transforms/buffer-hoisting.mlir
@@ -459,6 +459,46 @@ func @nested_region_control_flow_div_nested(
 
 // -----
 
+// Test Case: deeply nested region control flow with a nested buffer allocation
+// that has dependency within a nested region should not be moved outside of
+// this region.
+
+// CHECK-LABEL: func @nested_region_control_flow_div_nested_dependencies
+func @nested_region_control_flow_div_nested_dependencies(
+  %arg0: i32,
+  %arg1: i1,
+  %arg2: index) -> memref<?x?xf32> {
+  %0 = scf.if %arg1 -> (memref<?x?xf32>) {
+    %1 = constant 1 : i32
+    %2 = addi %arg0, %1 : i32
+    %3 = index_cast %2 : i32 to index
+    %4 = alloc(%arg2, %3) : memref<?x?xf32>
+    scf.yield %4 : memref<?x?xf32>
+  } else {
+    %1 = constant 2 : i32
+    %2 = addi %arg0, %1 : i32
+    %3 = index_cast %2 : i32 to index
+    %4 = alloc(%arg2, %3) : memref<?x?xf32>
+    scf.yield %4 : memref<?x?xf32>
+  }
+  return %0 : memref<?x?xf32>
+}
+
+//      CHECK: (%[[ARG0:.*]]: {{.*}}
+// CHECK-NEXT: %{{.*}} = scf.if
+// CHECK-NEXT: %{{.*}} = constant
+// CHECK-NEXT: %{{.*}} = addi
+// CHECK-NEXT: %[[FUNC:.*]] = index_cast
+// CHECK-NEXT: alloc(%arg2, %[[FUNC]])
+// CHECK-NEXT: scf.yield
+// CHECK-NEXT: } else {
+// CHECK-NEXT: %{{.*}} = constant
+// CHECK-NEXT: %{{.*}} = addi
+// CHECK-NEXT: %[[FUNC:.*]] = index_cast
+// CHECK-NEXT: alloc(%arg2, %[[FUNC]])
+
+// -----
+
 // Test Case: nested region control flow within a region interface.
 // The alloc positions of %0 does not need to be changed.