[flang-commits] [flang] [mlir] [flang][OpenMP] Lower target in_reduction for host fallback (PR #199967)

[Sergio Afonso via flang-commits]

================
@@ -8222,6 +8226,61 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
   bool isOffloadEntry =
       isTargetDevice || !ompBuilder->Config.TargetTriples.empty();
 
+  // Validate and resolve in_reduction clauses on omp.target. We currently
+  // only support the non-offload host-fallback path: the per-task private
+  // pointer is obtained by calling __kmpc_task_reduction_get_th_data inside
+  // the to-be-outlined target task body. Threading that pointer through the
+  // device kernel argument list is left as follow-up work.
+  SmallVector<llvm::Value *> inRedOrigPtrs;
+  if (!targetOp.getInReductionVars().empty()) {
+    if (isTargetDevice || isOffloadEntry)
+      return opInst.emitError(
+          "not yet implemented: in_reduction clause on omp.target with "
+          "offload / target-device compilation");
+    if (auto inRedSyms = targetOp.getInReductionSyms()) {
+      for (auto sym : inRedSyms->getAsRange<SymbolRefAttr>()) {
+        auto decl =
+            SymbolTable::lookupNearestSymbolFrom<omp::DeclareReductionOp>(
+                targetOp, sym);
+        if (!decl)
+          return targetOp.emitError()
+                 << "failed to resolve in_reduction declare_reduction symbol "
+                 << sym.getRootReference() << " on omp.target";
+        if (decl.getInitializerRegion().front().getNumArguments() != 1)
+          return targetOp.emitError()
+                 << "not yet implemented: in_reduction with two-argument "
+                    "initializer on omp.target";
+        if (!decl.getCleanupRegion().empty())
+          return targetOp.emitError()
+                 << "not yet implemented: in_reduction with cleanup region "
+                    "on omp.target";
+        // The reduction combiner region is intentionally not required here:
+        // the in_reduction lowering on omp.target only locates the per-task
+        // private storage via __kmpc_task_reduction_get_th_data. The combiner
+        // is owned by the enclosing taskgroup's task_reduction registration.
+      }
+    }
+    // Each in_reduction variable must also be captured by the target via a
+    // map_entries entry referring to the same outer SSA value. OMPIRBuilder
+    // outlines the target body and only rewires uses of values that enter
+    // the kernel through the map-derived input set. The runtime call below
+    // uses that same outer SSA value as its `orig` argument, so without a
+    // matching map entry the outlined kernel would reference a value defined
+    // in the host function and fail IR verification.
+    llvm::SmallPtrSet<Value, 4> mappedVarPtrs;
+    for (Value mapV : targetOp.getMapVars())
+      if (auto mapInfo = mapV.getDefiningOp<omp::MapInfoOp>())
+        mappedVarPtrs.insert(mapInfo.getVarPtr());
+    inRedOrigPtrs.reserve(targetOp.getInReductionVars().size());
+    for (Value v : targetOp.getInReductionVars()) {
+      if (!mappedVarPtrs.contains(v))
+        return targetOp.emitError()
+               << "not yet implemented: in_reduction variable on omp.target "
+                  "must also be captured by a matching map_entries entry";
----------------
skatrak wrote:

I think the fact that an `in_reduction` operand must also be a `map` argument is something that ought to be verified regardless of which one of the two representation alternatives we choose. Either way, I'm happier with the block argument-less version that's currently implemented in this patch.

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