diff --git a/flang/lib/Optimizer/Transforms/LoopVersioning.cpp b/flang/lib/Optimizer/Transforms/LoopVersioning.cpp index 4d3ea51ae1a5f..d030a3b3831ec 100644 --- a/flang/lib/Optimizer/Transforms/LoopVersioning.cpp +++ b/flang/lib/Optimizer/Transforms/LoopVersioning.cpp @@ -51,6 +51,7 @@ #include "flang/Optimizer/Dialect/Support/KindMapping.h" #include "flang/Optimizer/Transforms/Passes.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/IR/Dominance.h" #include "mlir/IR/Matchers.h" #include "mlir/IR/TypeUtilities.h" #include "mlir/Pass/Pass.h" @@ -292,6 +293,8 @@ void LoopVersioningPass::runOnOperation() { // immediately nested in a loop. llvm::DenseMap argsInLoops; + auto &domInfo = getAnalysis(); + // Traverse the loops in post-order and see // if those arguments are used inside any loop. func.walk([&](fir::DoLoopOp loop) { @@ -309,12 +312,15 @@ void LoopVersioningPass::runOnOperation() { for (auto a : argsOfInterest) { if (a.arg == normaliseVal(operand)) { // Use the reboxed value, not the block arg when re-creating the loop. - // TODO: should we check that the operand dominates the loop? - // If this might be a case, we should record such operands in - // argsInLoop.cannotTransform, so that they disable the transformation - // for the parent loops as well. a.arg = operand; + // Check that the operand dominates the loop? + // If this is the case, record such operands in argsInLoop.cannot- + // Transform, so that they disable the transformation for the parent + /// loops as well. + if (!domInfo.dominates(a.arg, loop)) + argsInLoop.cannotTransform.insert(a.arg); + // No support currently for sliced arrays. // This means that we cannot transform properly // instructions referencing a.arg in the whole loop diff --git a/flang/test/Transforms/loop-versioning.fir b/flang/test/Transforms/loop-versioning.fir index f2768d7325f74..6313bc2ac36a7 100644 --- a/flang/test/Transforms/loop-versioning.fir +++ b/flang/test/Transforms/loop-versioning.fir @@ -1434,4 +1434,48 @@ func.func @_QPtest_loop_nest(%arg0: !fir.box> {fir.bindc_name // CHECK: fir.if // CHECK-NOT: fir.if + +//----- + +// Check that a non-dominating operand isn't causing a problem. +// Just check it compiles, and doesn't version this loop. +func.func @sum1drebox(%arg0: !fir.box> {fir.bindc_name = "a"}, %arg1: !fir.ref {fir.bindc_name = "n"}) { + %decl = fir.declare %arg0 {uniq_name = "a"} : (!fir.box>) -> !fir.box> + %0 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QMmoduleFsum1dEi"} + %1 = fir.alloca f64 {bindc_name = "sum", uniq_name = "_QMmoduleFsum1dEsum"} + %cst = arith.constant 0.000000e+00 : f64 + fir.store %cst to %1 : !fir.ref + %c1_i32 = arith.constant 1 : i32 + %2 = fir.convert %c1_i32 : (i32) -> index + %3 = fir.load %arg1 : !fir.ref + %4 = fir.convert %3 : (i32) -> index + %c1 = arith.constant 1 : index + %5 = fir.convert %2 : (index) -> i32 + %6:2 = fir.do_loop %arg2 = %2 to %4 step %c1 iter_args(%arg3 = %5) -> (index, i32) { + // rebox is not dominating the loop. + %rebox = fir.rebox %decl : (!fir.box>) -> !fir.box> + fir.store %arg3 to %0 : !fir.ref + %7 = fir.load %1 : !fir.ref + %8 = fir.load %0 : !fir.ref + %9 = fir.convert %8 : (i32) -> i64 + %c1_i64 = arith.constant 1 : i64 + %10 = arith.subi %9, %c1_i64 : i64 + %11 = fir.coordinate_of %rebox, %10 : (!fir.box>, i64) -> !fir.ref + %12 = fir.load %11 : !fir.ref + %13 = arith.addf %7, %12 fastmath : f64 + fir.store %13 to %1 : !fir.ref + %14 = arith.addi %arg2, %c1 : index + %15 = fir.convert %c1 : (index) -> i32 + %16 = fir.load %0 : !fir.ref + %17 = arith.addi %16, %15 : i32 + fir.result %14, %17 : index, i32 + } + fir.store %6#1 to %0 : !fir.ref + return +} +// CHECK-LABEL: func @sum1drebox +// No versioning -> no if-operation. +// CHECK-NOT: fir.if + + } // End module