[Mlir-commits] [mlir] [mlir][OpenMP] Fix teams reduction assert with sibling distributes (PR #191475)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-openmp

Author: Chi-Chun, Chen (chichunchen)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/191475.diff


2 Files Affected:

- (modified) mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp (+20-17) 
- (added) mlir/test/Target/LLVMIR/openmp-teams-distribute-reduction-multiple.mlir (+44) 


``````````diff
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-multiple.mlir b/mlir/test/Target/LLVMIR/openmp-teams-distribute-reduction-multiple.mlir
new file mode 100644
index 0000000000000..a3fd2355c8efc
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-teams-distribute-reduction-multiple.mlir
@@ -0,0 +1,44 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// Regression test for a teams region with multiple nested distribute ops where
+// only one of them captures the teams reduction arguments.
+
+omp.declare_reduction @add_reduction_i32 : i32 init {
+^bb0(%arg0: i32):
+  %0 = llvm.mlir.constant(0 : i32) : i32
+  omp.yield(%0 : i32)
+} combiner {
+^bb0(%arg0: i32, %arg1: i32):
+  %0 = llvm.add %arg0, %arg1 : i32
+  omp.yield(%0 : i32)
+}
+
+llvm.func @teams_reduction_captured_by_one_distribute_only_() attributes {fir.internal_name = "_QPteams_reduction_captured_by_one_distribute_only", 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_captured_by_one_distribute_only_(
+// CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams
+// CHECK: call i32 @__kmpc_reduce
+// CHECK: call void @__kmpc_end_reduce

``````````

</details>


https://github.com/llvm/llvm-project/pull/191475