[flang][OpenMP] Lower target in_reduction for host fallback

Teach Flang lowering and MLIR OpenMP translation to carry
in_reduction through omp.target for the host-fallback path.

The translation looks up task reduction-private storage with
__kmpc_task_reduction_get_th_data and binds the target region's
in_reduction block argument to that private pointer, so uses inside the
region do not keep referring to the original variable.

The patch also preserves in_reduction operands in the TargetOp builder
path and ensures target in_reduction list items are mapped into the
target region when needed.

The device/offload-entry path remains diagnosed as not yet implemented.

Author: Saieiei

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -433,13 +433,16 @@ static void bindEntryBlockArgs(lower::AbstractConverter &converter,
               .first);
   };
 
-  // Process in clause name alphabetical order to match block arguments order.
   // Do not bind host_eval variables because they cannot be used inside of the
   // corresponding region, except for very specific cases handled separately.
+  // Bind map before in_reduction so that for target in_reduction list items
+  // (which are also implicitly mapped), the in_reduction binding wins and
+  // in-body references use the reduction-private block argument, not the
+  // mapped/original address.
   bindMapLike(args.hasDeviceAddr.objects, op.getHasDeviceAddrBlockArgs());
+  bindMapLike(args.map.objects, op.getMapBlockArgs());
   bindPrivateLike(args.inReduction.objects, args.inReduction.vars,
                   op.getInReductionBlockArgs());
-  bindMapLike(args.map.objects, op.getMapBlockArgs());
   bindPrivateLike(args.priv.objects, args.priv.vars, op.getPrivateBlockArgs());
   bindPrivateLike(args.reduction.objects, args.reduction.vars,
                   op.getReductionBlockArgs());
@@ -1873,6 +1876,7 @@ genTargetClauses(lower::AbstractConverter &converter,
                  mlir::omp::TargetOperands &clauseOps,
                  DefaultMapsTy &defaultMaps,
                  llvm::SmallVectorImpl<Object> &hasDeviceAddrObjects,
+                 llvm::SmallVectorImpl<Object> &inReductionObjects,
                  llvm::SmallVectorImpl<Object> &isDevicePtrObjects,
                  llvm::SmallVectorImpl<Object> &mapObjects) {
   ClauseProcessor cp(converter, semaCtx, clauses);
@@ -1887,13 +1891,14 @@ genTargetClauses(lower::AbstractConverter &converter,
     hostEvalInfo->collectValues(clauseOps.hostEvalVars);
   }
   cp.processIf(llvm::omp::Directive::OMPD_target, clauseOps);
+  cp.processInReduction(loc, clauseOps, inReductionObjects);
   cp.processIsDevicePtr(stmtCtx, clauseOps, isDevicePtrObjects);
   cp.processMap(loc, stmtCtx, clauseOps, llvm::omp::Directive::OMPD_unknown,
                 &mapObjects);
   cp.processNowait(clauseOps);
   cp.processThreadLimit(stmtCtx, clauseOps);
 
-  cp.processTODO<clause::Allocate, clause::InReduction, clause::UsesAllocators>(
+  cp.processTODO<clause::Allocate, clause::UsesAllocators>(
       loc, llvm::omp::Directive::OMPD_target);
 
   // `target private(..)` is only supported in delayed privatization mode.
@@ -2932,10 +2937,10 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
   mlir::omp::TargetOperands clauseOps;
   DefaultMapsTy defaultMaps;
   llvm::SmallVector<Object> mapObjects, hasDeviceAddrObjects,
-      isDevicePtrObjects;
+      inReductionObjects, isDevicePtrObjects;
   genTargetClauses(converter, semaCtx, symTable, stmtCtx, eval, item->clauses,
                    loc, clauseOps, defaultMaps, hasDeviceAddrObjects,
-                   isDevicePtrObjects, mapObjects);
+                   inReductionObjects, isDevicePtrObjects, mapObjects);
 
   if (!isDevicePtrObjects.empty()) {
     // is_device_ptr maps get duplicated so the clause and synthesized
@@ -3108,6 +3113,58 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
           Object{const_cast<semantics::Symbol *>(&sym), std::nullopt});
     }
   };
+  // OpenMP requires `in_reduction` list items on `target` to be implicitly
+  // data-mapped. The MLIR -> LLVM IR translation passes the mapped pointer
+  // as the `orig` argument of `__kmpc_task_reduction_get_th_data`, so the
+  // map must be address-preserving regardless of the scalar default capture
+  // (which would otherwise be ByCopy for small scalars and break the
+  // runtime lookup against the enclosing taskgroup's task_reduction
+  // descriptor). Emit these maps before the generic implicit-map walk so
+  // that walk treats the symbols as already mapped via
+  // `isDuplicateMappedSymbol` and does not downgrade them to ByCopy.
+  auto captureInReductionImplicitMap = [&](const semantics::Symbol &sym) {
+    if (sym.owner().IsDerivedType())
+      return;
+    if (!converter.getSymbolAddress(sym))
+      return;
+    if (isDuplicateMappedSymbol(sym, dsp.getAllSymbolsToPrivatize(),
+                                hasDeviceAddrObjects, mapObjects,
+                                isDevicePtrObjects))
+      return;
+    if (const auto *details =
+            sym.template detailsIf<semantics::HostAssocDetails>())
+      converter.copySymbolBinding(details->symbol(), sym);
+    std::stringstream name;
+    fir::ExtendedValue dataExv = converter.getSymbolExtendedValue(sym);
+    name << sym.name().ToString();
+    fir::factory::AddrAndBoundsInfo info =
+        Fortran::lower::getDataOperandBaseAddr(converter, firOpBuilder,
+                                               sym.GetUltimate(),
+                                               converter.getCurrentLocation());
+    llvm::SmallVector<mlir::Value> bounds =
+        fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
+                                           mlir::omp::MapBoundsType>(
+            firOpBuilder, info, dataExv,
+            semantics::IsAssumedSizeArray(sym.GetUltimate()),
+            converter.getCurrentLocation());
+    mlir::Value baseOp = info.rawInput;
+    mlir::omp::ClauseMapFlags flags = mlir::omp::ClauseMapFlags::implicit |
+                                      mlir::omp::ClauseMapFlags::to |
+                                      mlir::omp::ClauseMapFlags::from;
+    mlir::Value mapOp = createMapInfoOp(
+        firOpBuilder, converter.getCurrentLocation(), baseOp,
+        /*varPtrPtr=*/mlir::Value{}, name.str(), bounds, /*members=*/{},
+        /*membersIndex=*/mlir::ArrayAttr{}, flags,
+        mlir::omp::VariableCaptureKind::ByRef, baseOp.getType(),
+        /*partialMap=*/false, /*mapperId=*/mlir::FlatSymbolRefAttr{});
+    clauseOps.mapVars.push_back(mapOp);
+    mapObjects.push_back(
+        Object{const_cast<semantics::Symbol *>(&sym), std::nullopt});
+  };
+  for (const Object &object : inReductionObjects)
+    if (const semantics::Symbol *sym = object.sym())
+      captureInReductionImplicitMap(*sym);
+
   lower::pft::visitAllSymbols(eval, captureImplicitMap);
 
   auto targetOp = mlir::omp::TargetOp::create(firOpBuilder, loc, clauseOps);
@@ -3120,7 +3177,8 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
   args.hasDeviceAddr.objects = hasDeviceAddrObjects;
   args.hasDeviceAddr.vars = hasDeviceAddrBaseValues;
   args.hostEvalVars = clauseOps.hostEvalVars;
-  // TODO: Add in_reduction syms and vars.
+  args.inReduction.objects = inReductionObjects;
+  args.inReduction.vars = clauseOps.inReductionVars;
   args.map.objects = mapObjects;
   args.map.vars = mapBaseValues;
   args.priv.objects = makeObjects(dsp.getDelayedPrivSymbols());

flang/test/Lower/OpenMP/Todo/target-inreduction.f90

@@ -0,0 +1,27 @@
+! RUN: bbc -emit-hlfir -fopenmp -fopenmp-version=50 -o - %s 2>&1 | FileCheck %s
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=50 -o - %s 2>&1 | FileCheck %s
+
+! Per the OpenMP spec, an in_reduction list item on a target construct is
+! implicitly data-mapped. The lowering must not rely on the variable being
+! referenced inside the target body to discover that map: here `i` only
+! appears in the in_reduction clause and is never read or written inside
+! the region. Verify that an omp.map.info for `i` is still emitted and
+! flows into the omp.target's map_entries.
+
+!CHECK-LABEL: func.func @_QPomp_target_in_reduction_unused()
+!CHECK:       %[[IDECL:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "_QFomp_target_in_reduction_unusedEi"}
+!CHECK:       %[[IMAP:.*]] = omp.map.info var_ptr(%[[IDECL]]#1 : !fir.ref<i32>, i32) map_clauses(implicit, tofrom) capture(ByRef) -> !fir.ref<i32> {name = "i"}
+!CHECK:       omp.target in_reduction(@{{[^ ]+}} %[[IDECL]]#0 -> %{{[^ ]+}} : !fir.ref<i32>)
+!CHECK-SAME:    map_entries(%[[IMAP]] -> %{{[^ ]+}} : !fir.ref<i32>)
+
+subroutine omp_target_in_reduction_unused()
+  interface
+    subroutine sub()
+    end subroutine
+  end interface
+  integer i
+  i = 0
+  !$omp target in_reduction(+:i)
+  call sub()
+  !$omp end target
+end subroutine omp_target_in_reduction_unused

flang/test/Lower/OpenMP/target-inreduction-unused.f90

@@ -0,0 +1,28 @@
+! RUN: bbc -emit-hlfir -fopenmp -fopenmp-version=50 -o - %s 2>&1 | FileCheck %s
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=50 -o - %s 2>&1 | FileCheck %s
+
+! Verify that in_reduction on a target construct is lowered to an
+! omp.target with both an in_reduction clause and an implicit map_entries
+! entry for the same variable. The implicit map captures the original
+! pointer into the target region so the MLIR -> LLVM IR translation can
+! pass it to __kmpc_task_reduction_get_th_data.
+
+!CHECK-LABEL: omp.declare_reduction
+!CHECK-SAME:  @[[RED_I32_NAME:.*]] : i32 init {
+
+!CHECK-LABEL: func.func @_QPomp_target_in_reduction()
+!CHECK:       %[[IDECL:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "_QFomp_target_in_reductionEi"}
+!CHECK:       %[[IMAP:.*]] = omp.map.info var_ptr(%[[IDECL]]#1 : !fir.ref<i32>, i32) map_clauses(implicit, tofrom) capture(ByRef) -> !fir.ref<i32> {name = "i"}
+!CHECK:       omp.target in_reduction(@[[RED_I32_NAME]] %[[IDECL]]#0 -> %[[INARG:[^ ]+]] : !fir.ref<i32>)
+!CHECK-SAME:    map_entries(%[[IMAP]] -> %{{[^ ]+}} : !fir.ref<i32>)
+!CHECK:         hlfir.declare %[[INARG]]
+!CHECK:         omp.terminator
+!CHECK:       }
+
+subroutine omp_target_in_reduction()
+  integer i
+  i = 0
+  !$omp target in_reduction(+:i)
+  i = i + 1
+  !$omp end target
+end subroutine omp_target_in_reduction

flang/test/Lower/OpenMP/target-inreduction.f90

@@ -2545,18 +2545,18 @@ LogicalResult TargetUpdateOp::verify() {
 void TargetOp::build(OpBuilder &builder, OperationState &state,
                      const TargetOperands &clauses) {
   MLIRContext *ctx = builder.getContext();
-  // TODO Store clauses in op: allocateVars, allocatorVars, inReductionVars,
-  // inReductionByref, inReductionSyms.
+  // TODO Store clauses in op: allocateVars, allocatorVars.
   TargetOp::build(
       builder, state, /*allocate_vars=*/{}, /*allocator_vars=*/{}, clauses.bare,
       makeArrayAttr(ctx, clauses.dependKinds), clauses.dependVars,
       makeArrayAttr(ctx, clauses.dependIteratedKinds), clauses.dependIterated,
       clauses.device, clauses.dynGroupprivateAccessGroup,
       clauses.dynGroupprivateFallback, clauses.dynGroupprivateSize,
       clauses.hasDeviceAddrVars, clauses.hostEvalVars, clauses.ifExpr,
-      /*in_reduction_vars=*/{}, /*in_reduction_byref=*/nullptr,
-      /*in_reduction_syms=*/nullptr, clauses.isDevicePtrVars, clauses.mapVars,
-      clauses.nowait, clauses.privateVars,
+      clauses.inReductionVars,
+      makeDenseBoolArrayAttr(ctx, clauses.inReductionByref),
+      makeArrayAttr(ctx, clauses.inReductionSyms), clauses.isDevicePtrVars,
+      clauses.mapVars, clauses.nowait, clauses.privateVars,
       makeArrayAttr(ctx, clauses.privateSyms), clauses.privateNeedsBarrier,
       clauses.threadLimitVars,
       /*private_maps=*/nullptr);
@@ -2583,6 +2583,11 @@ LogicalResult TargetOp::verify() {
   if (failed(verifyPrivateVarList(*this)))
     return failure();
 
+  if (failed(verifyReductionVarList(*this, getInReductionSyms(),
+                                    getInReductionVars(),
+                                    getInReductionByref())))
+    return failure();
+
   return verifyPrivateVarsMapping(*this);
 }
 

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

@@ -490,7 +490,11 @@ static LogicalResult checkImplementationStatus(Operation &op) {
       .Case([&](omp::TargetOp op) {
         checkAllocate(op, result);
         checkBare(op, result);
-        checkInReduction(op, result);
+        // in_reduction(byref(...)) on target is not implemented yet. Other
+        // unsupported in_reduction shapes (cleanup region, two-argument
+        // initializer, missing combiner) and the device-side / offload-entry
+        // cases are diagnosed inline in convertOmpTarget.
+        checkInReductionByref(op, result);
         checkThreadLimit(op, result);
       })
       .Default([](Operation &) {
@@ -8208,6 +8212,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";
+      inRedOrigPtrs.push_back(moduleTranslation.lookupValue(v));
+    }
+  }
+
   // For some private variables, the MapsForPrivatizedVariablesPass
   // creates MapInfoOp instances. Go through the private variables and
   // the mapped variables so that during codegeneration we are able
@@ -8320,6 +8379,36 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
             targetOp.getPrivateNeedsBarrier(), &mappedPrivateVars)))
       return llvm::make_error<PreviouslyReportedError>();
 
+    // Map in_reduction block arguments to the per-task private storage
+    // returned by __kmpc_task_reduction_get_th_data. The lookup must run
+    // inside the target task body so the gtid corresponds to the executing
+    // thread. The descriptor argument is NULL: the runtime walks enclosing
+    // taskgroups to locate the matching task_reduction registration for
+    // `origPtr`. Mirrors the in_reduction handling on omp.taskloop.context.
+    ArrayRef<BlockArgument> inRedBlockArgs = argIface.getInReductionBlockArgs();
+    if (!inRedBlockArgs.empty()) {
+      llvm::OpenMPIRBuilder &ompB = *moduleTranslation.getOpenMPBuilder();
+      llvm::Module *m = moduleTranslation.getLLVMModule();
+      llvm::LLVMContext &llvmCtx = m->getContext();
+      uint32_t srcLocSize;
+      llvm::Constant *srcLocStr = ompB.getOrCreateDefaultSrcLocStr(srcLocSize);
+      llvm::Value *bodyIdent = ompB.getOrCreateIdent(srcLocStr, srcLocSize);
+      llvm::Function *gtidFn = ompB.getOrCreateRuntimeFunctionPtr(
+          llvm::omp::OMPRTL___kmpc_global_thread_num);
+      llvm::Value *bodyGtid =
+          builder.CreateCall(gtidFn, {bodyIdent}, "omp_global_thread_num");
+      llvm::FunctionCallee getThData = ompB.getOrCreateRuntimeFunction(
+          *m, llvm::omp::OMPRTL___kmpc_task_reduction_get_th_data);
+      llvm::Type *ptrTy = llvm::PointerType::getUnqual(llvmCtx);
+      llvm::Value *nullDesc = llvm::ConstantPointerNull::get(ptrTy);
+      for (auto [blockArg, origPtr] :
+           llvm::zip_equal(inRedBlockArgs, inRedOrigPtrs)) {
+        llvm::Value *priv = builder.CreateCall(
+            getThData, {bodyGtid, nullDesc, origPtr}, "omp.inred.priv");
+        moduleTranslation.mapValue(blockArg, priv);
+      }
+    }
+
     LLVM::ModuleTranslation::SaveStack<OpenMPAllocStackFrame> frame(
         moduleTranslation, allocaIP, deallocBlocks);
     llvm::Expected<llvm::BasicBlock *> exitBlock = convertOmpOpRegions(