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

[Sairudra More]

================
@@ -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();
----------------
Saieiei wrote:

Done!

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