diff --git a/mlir/lib/Transforms/BufferOptimizations.cpp b/mlir/lib/Transforms/BufferOptimizations.cpp --- a/mlir/lib/Transforms/BufferOptimizations.cpp +++ b/mlir/lib/Transforms/BufferOptimizations.cpp @@ -142,13 +142,13 @@ // 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 @@ explicit PromoteBuffersToStackPass(std::function 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 --- a/mlir/test/Transforms/buffer-hoisting.mlir +++ b/mlir/test/Transforms/buffer-hoisting.mlir @@ -459,6 +459,46 @@ // ----- +// 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 { + %0 = scf.if %arg1 -> (memref) { + %1 = constant 1 : i32 + %2 = addi %arg0, %1 : i32 + %3 = index_cast %2 : i32 to index + %4 = alloc(%arg2, %3) : memref + scf.yield %4 : memref + } else { + %1 = constant 2 : i32 + %2 = addi %arg0, %1 : i32 + %3 = index_cast %2 : i32 to index + %4 = alloc(%arg2, %3) : memref + scf.yield %4 : memref + } + return %0 : memref +} + +// 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.