[clang] [llvm] [mlir] [MLIR][OpenMP] Add codegen for teams reductions (PR #133310)

[Sergio Afonso via llvm-commits]

================
@@ -1714,6 +1727,42 @@ convertOmpSingle(omp::SingleOp &singleOp, llvm::IRBuilderBase &builder,
   return success();
 }
 
+static bool teamsReductionContainedInDistribute(omp::TeamsOp teamsOp) {
+  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();
+      // Ignore debug uses.
+      if (mlir::isa<LLVM::DbgDeclareOp, LLVM::DbgValueOp>(useOp)) {
+        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 we are going to use distribute reduction then remove any debug uses of
+  // the reduction parameters in teamsOp. Otherwise they will be left without
+  // any mapped value in moduleTranslation and will eventually error out.
+  for (auto use : debugUses)
+    use->erase();
----------------
skatrak wrote:

@abidh, can you take a look at whether this is safe/the right thing to do?

It seems to me, though I'm not familiar with debug information handling, that another way to deal with this would be to map these `omp.teams` entry block arguments with the corresponding LLVM values defined while handling the reduction inside of the nested `omp.distribute`. Then, we would just have to sink these debug ops into `omp.loop_nest` rather than deleting them and we would be able to keep that information.

I don't actually know if we'd be losing any information by deleting these ops, this is just an idea in case we wanted to keep them.

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