[Mlir-commits] [mlir] [mlir][OpenMP] Translate reductions on taskloop (PR #199670)

[Tom Eccles]

================
@@ -3417,6 +3443,90 @@ convertOmpTaskloopContextOp(omp::TaskloopContextOp contextOp,
   // Set up inserttion point for call to createTaskloop()
   builder.SetInsertPoint(taskloopStartBlock);
 
+  // Resolve and validate reduction / in_reduction declarations. Only the
+  // non-byref, single-init-arg, no-cleanup form is supported in this first
+  // cut; richer shapes have been rejected by checkImplementationStatus
+  // (byref) or are rejected here.
+  auto resolveRedDecls =
+      [&](std::optional<ArrayAttr> syms, StringRef clauseName,
+          SmallVectorImpl<omp::DeclareReductionOp> &out) -> LogicalResult {
+    if (!syms)
+      return success();
+    out.reserve(syms->size());
+    for (auto sym : syms->getAsRange<SymbolRefAttr>()) {
+      auto decl = SymbolTable::lookupNearestSymbolFrom<omp::DeclareReductionOp>(
+          contextOp, sym);
+      if (!decl)
+        return contextOp.emitError()
+               << "failed to resolve " << clauseName
+               << " declare_reduction symbol " << sym.getRootReference()
+               << " in omp.taskloop.context";
+      if (decl.getInitializerRegion().front().getNumArguments() != 1)
+        return contextOp.emitError()
+               << "not yet implemented: " << clauseName
+               << " with two-argument initializer in omp.taskloop.context";
+      if (!decl.getCleanupRegion().empty())
+        return contextOp.emitError()
+               << "not yet implemented: " << clauseName
+               << " with cleanup region in omp.taskloop.context";
+      if (decl.getReductionRegion().empty())
+        return contextOp.emitError()
+               << clauseName
+               << " declare_reduction is missing a combiner region";
+      out.push_back(decl);
+    }
+    return success();
+  };
+
+  SmallVector<omp::DeclareReductionOp> redDecls;
+  if (failed(
+          resolveRedDecls(contextOp.getReductionSyms(), "reduction", redDecls)))
+    return failure();
+  SmallVector<omp::DeclareReductionOp> inRedDecls;
+  if (failed(resolveRedDecls(contextOp.getInReductionSyms(), "in_reduction",
+                             inRedDecls)))
+    return failure();
+
+  // The op verifier rejects nogroup + reduction, so no check is needed here.
+
+  SmallVector<llvm::Value *> redOrigPtrs;
+  redOrigPtrs.reserve(redDecls.size());
+  for (Value v : contextOp.getReductionVars())
+    redOrigPtrs.push_back(moduleTranslation.lookupValue(v));
+  SmallVector<llvm::Value *> inRedOrigPtrs;
+  inRedOrigPtrs.reserve(inRedDecls.size());
+  for (Value v : contextOp.getInReductionVars())
+    inRedOrigPtrs.push_back(moduleTranslation.lookupValue(v));
+
+  llvm::OpenMPIRBuilder &ompBuilderRef = *moduleTranslation.getOpenMPBuilder();
+  llvm::Module *llvmModuleForRed = moduleTranslation.getLLVMModule();
+
+  // If we have task_reduction items, we must emit our own implicit
+  // __kmpc_taskgroup so that the descriptor returned by __kmpc_taskred_init
+  // is associated with that taskgroup. We then force NoGroup=true so that
+  // OpenMPIRBuilder::createTaskloop does not emit a second taskgroup.
----------------
tblah wrote:

In this case, I think it would be cleaner to just create the implicit taskgroup in OpenMPToLLVMIRTranslation whenever one will be needed and always pass nogroup to OpenMPIRBuilder. This could be done in another commit in preparation for this one.

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