[Mlir-commits] [mlir] [mlir][OpenMP] Fix teams reduction assert with sibling distributes (PR #191475)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Apr 14 10:04:27 PDT 2026
https://github.com/chichunchen updated https://github.com/llvm/llvm-project/pull/191475
>From bb19f599b8357ecd782c1f08d808674cdb6da2a9 Mon Sep 17 00:00:00 2001
From: "Chi Chun, Chen" <chichun.chen at hpe.com>
Date: Fri, 10 Apr 2026 02:53:07 -0500
Subject: [PATCH 1/2] [mlir][OpenMP] Fix teams reduction assert with sibling
distributes
Avoid double-mapping teams reduction block arguments when multiple
sibling `omp.distribute` ops appear under the same `omp.teams`.
Replace `teamsReductionContainedInDistribute` with
`getDistributeCapturingTeamsReduction`, which returns the unique
`omp::DistributeOp` containing all non-debug uses of the teams reduction
block arguments, or null if no such distribute exists. In that case,
reduction setup falls back to the teams op.
This ensures only the owning distribute handles reduction setup and
prevents `moduleTranslation.mapValue()` from being called multiple times
for the same block arguments.
Fixes https://github.com/llvm/llvm-project/issues/191470
Assisted with Copilot and GPT-5.4
---
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 37 ++++++++++---------
.../openmp-teams-distribute-reduction.mlir | 36 ++++++++++++++++--
2 files changed, 53 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 2e15f4de4545d..cceee1a2a3143 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::advance();
+ });
+ if (walk.wasInterrupted() || !distOp)
+ return {};
+
auto iface =
llvm::cast<mlir::omp::BlockArgOpenMPOpInterface>(teamsOp.getOperation());
// Check that all uses of the reduction block arg has the same distribute op
// parent.
llvm::SmallVector<mlir::Operation *> 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<omp::DistributeOp>();
- // 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<omp::TeamsOp>();
bool doDistributeReduction =
- teamsOp ? teamsReductionContainedInDistribute(teamsOp) : false;
+ teamsOp && getDistributeCapturingTeamsReduction(teamsOp) == distributeOp;
DenseMap<Value, llvm::Value *> 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..a367520922ebe 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,35 @@ llvm.func @simple_teams_reduction_() attributes {fir.internal_name = "_QPsimple_
// Non atomic version
// CHECK: call void @__kmpc_end_reduce
-// Reduction function.
-// CHECK: define internal void @[[REDFUNC]]
-// CHECK: add i32
+// 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<all>, 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: call void @__kmpc_end_reduce
>From 9e2ef5765cf8aba9c9e116273119c14076d29d49 Mon Sep 17 00:00:00 2001
From: "Chi Chun, Chen" <chichun.chen at hpe.com>
Date: Tue, 14 Apr 2026 12:00:33 -0500
Subject: [PATCH 2/2] Skip contents of omp.distribute region and add back check
lines
---
.../LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp | 2 +-
.../test/Target/LLVMIR/openmp-teams-distribute-reduction.mlir | 4 ++++
2 files changed, 5 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index cceee1a2a3143..bef0b90db29d1 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -1990,7 +1990,7 @@ getDistributeCapturingTeamsReduction(omp::TeamsOp teamsOp) {
if (distOp)
return WalkResult::interrupt();
distOp = op;
- return WalkResult::advance();
+ return WalkResult::skip();
});
if (walk.wasInterrupted() || !distOp)
return {};
diff --git a/mlir/test/Target/LLVMIR/openmp-teams-distribute-reduction.mlir b/mlir/test/Target/LLVMIR/openmp-teams-distribute-reduction.mlir
index a367520922ebe..37f759f08cb18 100644
--- a/mlir/test/Target/LLVMIR/openmp-teams-distribute-reduction.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-teams-distribute-reduction.mlir
@@ -93,3 +93,7 @@ llvm.func @teams_reduction_multiple_distribute_() attributes {fir.internal_name
// CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams
// CHECK: call i32 @__kmpc_reduce
// CHECK: call void @__kmpc_end_reduce
+
+// Reduction function.
+// CHECK: define internal void @{{.+}}(ptr %0, ptr %1)
+// CHECK: add i32
More information about the Mlir-commits
mailing list