diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp index 2e15f4de4545d..bef0b90db29d1 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp @@ -1981,13 +1981,25 @@ convertOmpSingle(omp::SingleOp &singleOp, llvm::IRBuilderBase &builder, return success(); } -static bool teamsReductionContainedInDistribute(omp::TeamsOp teamsOp) { +static omp::DistributeOp +getDistributeCapturingTeamsReduction(omp::TeamsOp teamsOp) { + // Early return if we found more than one distribute op or if we can't find + // any distribute op in the teams region. + omp::DistributeOp distOp; + WalkResult walk = teamsOp.getRegion().walk([&](omp::DistributeOp op) { + if (distOp) + return WalkResult::interrupt(); + distOp = op; + return WalkResult::skip(); + }); + if (walk.wasInterrupted() || !distOp) + return {}; + auto iface = llvm::cast(teamsOp.getOperation()); // Check that all uses of the reduction block arg has the same distribute op // parent. llvm::SmallVector debugUses; - Operation *distOp = nullptr; for (auto ra : iface.getReductionBlockArgs()) for (auto &use : ra.getUses()) { auto *useOp = use.getOwner(); @@ -1996,17 +2008,8 @@ static bool teamsReductionContainedInDistribute(omp::TeamsOp teamsOp) { debugUses.push_back(useOp); continue; } - - auto currentDistOp = useOp->getParentOfType(); - // Use is not inside a distribute op - return false - if (!currentDistOp) - return false; - // Multiple distribute operations - return false - Operation *currentOp = currentDistOp.getOperation(); - if (distOp && (distOp != currentOp)) - return false; - - distOp = currentOp; + if (!distOp->isProperAncestor(useOp)) + return {}; } // If we are going to use distribute reduction then remove any debug uses of @@ -2014,7 +2017,7 @@ static bool teamsReductionContainedInDistribute(omp::TeamsOp teamsOp) { // any mapped value in moduleTranslation and will eventually error out. for (auto *use : debugUses) use->erase(); - return true; + return distOp; } // Convert an OpenMP Teams construct to LLVM IR using OpenMPIRBuilder @@ -2035,7 +2038,7 @@ convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder, // Only do teams reduction if there is no distribute op that captures the // reduction instead. - bool doTeamsReduction = !teamsReductionContainedInDistribute(op); + bool doTeamsReduction = !getDistributeCapturingTeamsReduction(op); if (doTeamsReduction) { isByRef = getIsByRef(op.getReductionByref()); @@ -6120,10 +6123,10 @@ convertOmpDistribute(Operation &opInst, llvm::IRBuilderBase &builder, return failure(); /// Process teams op reduction in distribute if the reduction is contained in - /// the distribute op. + /// this specific distribute op. omp::TeamsOp teamsOp = opInst.getParentOfType(); bool doDistributeReduction = - teamsOp ? teamsReductionContainedInDistribute(teamsOp) : false; + teamsOp && getDistributeCapturingTeamsReduction(teamsOp) == distributeOp; DenseMap reductionVariableMap; unsigned numReductionVars = teamsOp ? teamsOp.getNumReductionVars() : 0; diff --git a/mlir/test/Target/LLVMIR/openmp-teams-distribute-reduction.mlir b/mlir/test/Target/LLVMIR/openmp-teams-distribute-reduction.mlir index ac58e07e3ba37..066d40a020a2b 100644 --- a/mlir/test/Target/LLVMIR/openmp-teams-distribute-reduction.mlir +++ b/mlir/test/Target/LLVMIR/openmp-teams-distribute-reduction.mlir @@ -39,6 +39,7 @@ llvm.func @simple_teams_reduction_() attributes {fir.internal_name = "_QPsimple_ } llvm.return } +// CHECK-LABEL: define void @simple_teams_reduction_( // Call to outlined function // CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams // CHECK-SAME: @[[OUTLINED:[A-Za-z_.][A-Za-z0-9_.]*]] @@ -60,6 +61,40 @@ llvm.func @simple_teams_reduction_() attributes {fir.internal_name = "_QPsimple_ // Non atomic version // CHECK: call void @__kmpc_end_reduce +// Regression test: teams region with multiple sibling distribute ops where only +// one captures the teams reduction arguments. The teams op must handle the +// reduction itself to avoid a duplicate mapValue assertion. +llvm.func @teams_reduction_multiple_distribute_() attributes {fir.internal_name = "_QPteams_reduction_multiple_distribute", frame_pointer = #llvm.framePointerKind, target_cpu = "x86-64"} { + %0 = llvm.mlir.constant(1 : i64) : i64 + %1 = llvm.alloca %0 x i32 {bindc_name = "sum"} : (i64) -> !llvm.ptr + %2 = llvm.mlir.constant(1 : i32) : i32 + %3 = llvm.mlir.constant(0 : i32) : i32 + llvm.store %3, %1 : i32, !llvm.ptr + omp.teams reduction(@add_reduction_i32 %1 -> %arg0 : !llvm.ptr) { + omp.distribute { + omp.loop_nest (%arg1) : i32 = (%2) to (%2) inclusive step (%2) { + omp.yield + } + } + omp.distribute { + omp.loop_nest (%arg1) : i32 = (%2) to (%2) inclusive step (%2) { + %4 = llvm.load %arg0 : !llvm.ptr -> i32 + %5 = llvm.add %4, %arg1 : i32 + llvm.store %5, %arg0 : i32, !llvm.ptr + omp.yield + } + } + omp.terminator + } + llvm.return +} + +// CHECK-LABEL: define void @teams_reduction_multiple_distribute_( +// CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams +// CHECK: call i32 @__kmpc_reduce +// CHECK-SAME: @[[REDFUNC:[A-Za-z_.][A-Za-z0-9_.]*]] +// CHECK: call void @__kmpc_end_reduce + // Reduction function. // CHECK: define internal void @[[REDFUNC]] // CHECK: add i32