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

[dfki-jugr]

Summary

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

[herhut]

Thanks!

mlir/lib/Transforms/BufferOptimizations.cpp

for (BufferPlacementAllocs::AllocEntry &entry : allocs) {
       // Determine the common dominator block of all aliases.
       Block *dominatorBlock =
           findCommonDominator(allocValue, resultAliases, dominators);
       // Init the initial hoisting state.
       StateT state(&dominators, allocValue, allocValue.getParentBlock());
       // 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
       // an upper placement-block boundary. The idea is that placement block
       // cannot be moved any further upwards than the given upper bound.
       Block *placementBlock = findPlacementBlock(
           state, state.computeUpperBound(dominatorBlock, dependencyBlock));
       Operation *startOperation = BufferPlacementAllocs::getStartOperation(
PromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,
     this->maxAllocSizeInBytes = maxAllocSizeInBytes;
     this->bitwidthOfIndexType = bitwidthOfIndexType;
     this->maxRankOfAllocatedMemRef = maxRankOfAllocatedMemRef;
   }
 
   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,
                                    bitwidthOfIndexType,
                                    maxRankOfAllocatedMemRef);
       };
     }
     return success();

mlir/test/Transforms/buffer-hoisting.mlir

 }
 //      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
 // CHECK-NEXT: %[[ALLOC1:.*]] = alloc(%arg0, %arg1)
 // CHECK-NEXT: %[[ALLOC2:.*]] = alloc(%arg1, %arg1)
 // CHECK-NEXT: %{{.*}} = scf.if
 
 // -----
 
+// 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.
 
 // CHECK-LABEL: func @inner_region_control_flow
 func @inner_region_control_flow(%arg0 : index) -> memref<?x?xf32> {
   %0 = alloc(%arg0, %arg0) : memref<?x?xf32>
   %1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
     ^bb0(%arg1 : memref<?x?xf32>):