[Flang][OpenMP] Make implicitly captured scalars fully firstprivatized

Currently, we indicate to the runtime that implicit scalar captures are firstprivate (via map and
capture types), enough for the runtime trace to treat it as such, but we do not CodeGen the IR
in such a way that we can take full advantage of this aspect of the OpenMP specification.

This patch seeks to change that by applying the correct symbol flags (firstprivate/implicit) to the
implicitly captured scalars within target regions, which then triggers the delayed privitization code
generation for these symbols, bringing the code generation in-line with the explicit firstpriviate
clause. Currently, similarly to the delayed privitization I have sheltered this segment of code
behind the EnabledDelayedPrivitization flag, as without it, we'll trigger an compiler error for
firstprivate not being supported any time we implicitly capture a scalar and try to firstprivitize
it, in future when this flag is removed it can also be removed here. So, for now, you need to
enable this via providing the compiler the flag on compilation of any programs.

Author: agozillon

flang/lib/Lower/OpenMP/DataSharingProcessor.cpp

@@ -43,16 +43,42 @@ bool DataSharingProcessor::OMPConstructSymbolVisitor::isSymbolDefineBy(
                        [](const auto &functionParserNode) { return false; }});
 }
 
+static bool isConstructWithTopLevelTarget(lower::pft::Evaluation &eval) {
+  const auto *ompEval = eval.getIf<parser::OpenMPConstruct>();
+  if (ompEval) {
+    auto dir = parser::omp::GetOmpDirectiveName(ompEval).v;
+    switch (dir) {
+    case llvm::omp::Directive::OMPD_target:
+    case llvm::omp::Directive::OMPD_target_loop:
+    case llvm::omp::Directive::OMPD_target_parallel_do:
+    case llvm::omp::Directive::OMPD_target_parallel_do_simd:
+    case llvm::omp::Directive::OMPD_target_parallel_loop:
+    case llvm::omp::Directive::OMPD_target_teams_distribute:
+    case llvm::omp::Directive::OMPD_target_teams_distribute_parallel_do:
+    case llvm::omp::Directive::OMPD_target_teams_distribute_parallel_do_simd:
+    case llvm::omp::Directive::OMPD_target_teams_distribute_simd:
+    case llvm::omp::Directive::OMPD_target_teams_loop:
+    case llvm::omp::Directive::OMPD_target_simd:
+      return true;
+      break;
+    default:
+      return false;
+      break;
+    }
+  }
+  return false;
+}
+
 DataSharingProcessor::DataSharingProcessor(
     lower::AbstractConverter &converter, semantics::SemanticsContext &semaCtx,
     const List<Clause> &clauses, lower::pft::Evaluation &eval,
     bool shouldCollectPreDeterminedSymbols, bool useDelayedPrivatization,
-    lower::SymMap &symTable)
+    lower::SymMap &symTable, bool isTargetPrivitization)
     : converter(converter), semaCtx(semaCtx),
       firOpBuilder(converter.getFirOpBuilder()), clauses(clauses), eval(eval),
       shouldCollectPreDeterminedSymbols(shouldCollectPreDeterminedSymbols),
       useDelayedPrivatization(useDelayedPrivatization), symTable(symTable),
-      visitor(semaCtx) {
+      isTargetPrivitization(isTargetPrivitization), visitor(semaCtx) {
   eval.visit([&](const auto &functionParserNode) {
     parser::Walk(functionParserNode, visitor);
   });
@@ -62,10 +88,12 @@ DataSharingProcessor::DataSharingProcessor(lower::AbstractConverter &converter,
                                            semantics::SemanticsContext &semaCtx,
                                            lower::pft::Evaluation &eval,
                                            bool useDelayedPrivatization,
-                                           lower::SymMap &symTable)
+                                           lower::SymMap &symTable,
+                                           bool isTargetPrivitization)
     : DataSharingProcessor(converter, semaCtx, {}, eval,
                            /*shouldCollectPreDeterminedSymols=*/false,
-                           useDelayedPrivatization, symTable) {}
+                           useDelayedPrivatization, symTable,
+                           isTargetPrivitization) {}
 
 void DataSharingProcessor::processStep1(
     mlir::omp::PrivateClauseOps *clauseOps) {
@@ -552,8 +580,19 @@ void DataSharingProcessor::collectSymbols(
   };
 
   auto shouldCollectSymbol = [&](const semantics::Symbol *sym) {
-    if (collectImplicit)
+    if (collectImplicit) {
+      // If we're a combined construct with a target region, implicit
+      // firstprivate captures, should only belong to the target region
+      // and not be added/captured by later directives. Parallel regions
+      // will likely want the same captures to be shared and for SIMD it's
+      // illegal to have firstprivate clauses.
+      if (isConstructWithTopLevelTarget(eval) && !isTargetPrivitization &&
+          sym->test(semantics::Symbol::Flag::OmpFirstPrivate)) {
+        return false;
+      }
+
       return sym->test(semantics::Symbol::Flag::OmpImplicit);
+    }
 
     if (collectPreDetermined)
       return sym->test(semantics::Symbol::Flag::OmpPreDetermined);

flang/lib/Lower/OpenMP/DataSharingProcessor.h

@@ -93,6 +93,7 @@ class DataSharingProcessor {
   bool useDelayedPrivatization;
   llvm::SmallSet<const semantics::Symbol *, 16> mightHaveReadHostSym;
   lower::SymMap &symTable;
+  bool isTargetPrivitization;
   OMPConstructSymbolVisitor visitor;
 
   bool needBarrier();
@@ -130,12 +131,14 @@ class DataSharingProcessor {
                        const List<Clause> &clauses,
                        lower::pft::Evaluation &eval,
                        bool shouldCollectPreDeterminedSymbols,
-                       bool useDelayedPrivatization, lower::SymMap &symTable);
+                       bool useDelayedPrivatization, lower::SymMap &symTable,
+                       bool isTargetPrivitization = false);
 
   DataSharingProcessor(lower::AbstractConverter &converter,
                        semantics::SemanticsContext &semaCtx,
                        lower::pft::Evaluation &eval,
-                       bool useDelayedPrivatization, lower::SymMap &symTable);
+                       bool useDelayedPrivatization, lower::SymMap &symTable,
+                       bool isTargetPrivitization = false);
 
   // Privatisation is split into two steps.
   // Step1 performs cloning of all privatisation clauses and copying for

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -2465,6 +2465,36 @@ genSingleOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
       queue, item, clauseOps);
 }
 
+static bool isDuplicateMappedSymbol(
+    const semantics::Symbol &sym,
+    const llvm::SetVector<const semantics::Symbol *> &privatizedSyms,
+    const llvm::SmallVectorImpl<const semantics::Symbol *> &hasDevSyms,
+    const llvm::SmallVectorImpl<const semantics::Symbol *> &mappedSyms) {
+  llvm::SmallVector<const semantics::Symbol *> concatSyms;
+  concatSyms.reserve(privatizedSyms.size() + hasDevSyms.size() +
+                     mappedSyms.size());
+  concatSyms.append(privatizedSyms.begin(), privatizedSyms.end());
+  concatSyms.append(hasDevSyms.begin(), hasDevSyms.end());
+  concatSyms.append(mappedSyms.begin(), mappedSyms.end());
+
+  auto checkSymbol = [&](const semantics::Symbol &checkSym) {
+    if (llvm::is_contained(concatSyms, &checkSym))
+      return true;
+
+    return std::any_of(concatSyms.begin(), concatSyms.end(),
+                       [&](auto v) { return v->GetUltimate() == checkSym; });
+  };
+
+  if (checkSymbol(sym))
+    return true;
+
+  const auto *hostAssoc{sym.detailsIf<semantics::HostAssocDetails>()};
+  if (hostAssoc && checkSymbol(hostAssoc->symbol()))
+    return true;
+
+  return checkSymbol(sym.GetUltimate());
+}
+
 static mlir::omp::TargetOp
 genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
             lower::StatementContext &stmtCtx,
@@ -2491,7 +2521,8 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
   DataSharingProcessor dsp(converter, semaCtx, item->clauses, eval,
                            /*shouldCollectPreDeterminedSymbols=*/
                            lower::omp::isLastItemInQueue(item, queue),
-                           /*useDelayedPrivatization=*/true, symTable);
+                           /*useDelayedPrivatization=*/true, symTable,
+                           /*isTargetPrivitization=*/true);
   dsp.processStep1(&clauseOps);
 
   // 5.8.1 Implicit Data-Mapping Attribute Rules
@@ -2500,17 +2531,6 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
   // attribute clauses (neither data-sharing; e.g. `private`, nor `map`
   // clauses).
   auto captureImplicitMap = [&](const semantics::Symbol &sym) {
-    if (dsp.getAllSymbolsToPrivatize().contains(&sym))
-      return;
-
-    // Skip parameters/constants as they do not need to be mapped.
-    if (semantics::IsNamedConstant(sym))
-      return;
-
-    // These symbols are mapped individually in processHasDeviceAddr.
-    if (llvm::is_contained(hasDeviceAddrSyms, &sym))
-      return;
-
     // Structure component symbols don't have bindings, and can only be
     // explicitly mapped individually. If a member is captured implicitly
     // we map the entirety of the derived type when we find its symbol.
@@ -2532,7 +2552,12 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
     if (!converter.getSymbolAddress(sym))
       return;
 
-    if (!llvm::is_contained(mapSyms, &sym)) {
+    // Skip parameters/constants as they do not need to be mapped.
+    if (semantics::IsNamedConstant(sym))
+      return;
+
+    if (!isDuplicateMappedSymbol(sym, dsp.getAllSymbolsToPrivatize(),
+                                 hasDeviceAddrSyms, mapSyms)) {
       if (const auto *details =
               sym.template detailsIf<semantics::HostAssocDetails>())
         converter.copySymbolBinding(details->symbol(), sym);

flang/lib/Lower/Support/Utils.cpp

@@ -12,6 +12,7 @@
 
 #include "flang/Lower/Support/Utils.h"
 
+#include "flang/Common/idioms.h"
 #include "flang/Common/indirection.h"
 #include "flang/Lower/AbstractConverter.h"
 #include "flang/Lower/ConvertVariable.h"

flang/lib/Optimizer/OpenMP/MapsForPrivatizedSymbols.cpp

@@ -55,6 +55,20 @@ class MapsForPrivatizedSymbolsPass
 
   omp::MapInfoOp createMapInfo(Location loc, Value var,
                                fir::FirOpBuilder &builder) {
+    // Check if a value of type `type` can be passed to the kernel by value.
+    // All kernel parameters are of pointer type, so if the value can be
+    // represented inside of a pointer, then it can be passed by value.
+    auto isLiteralType = [&](mlir::Type type) {
+      const mlir::DataLayout &dl = builder.getDataLayout();
+      mlir::Type ptrTy = mlir::LLVM::LLVMPointerType::get(builder.getContext());
+      uint64_t ptrSize = dl.getTypeSize(ptrTy);
+      uint64_t ptrAlign = dl.getTypePreferredAlignment(ptrTy);
+
+      auto [size, align] = fir::getTypeSizeAndAlignmentOrCrash(
+          loc, type, dl, builder.getKindMap());
+      return size <= ptrSize && align <= ptrAlign;
+    };
+
     uint64_t mapTypeTo = static_cast<
         std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
         llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO);
@@ -94,14 +108,22 @@ class MapsForPrivatizedSymbolsPass
     if (needsBoundsOps(varPtr))
       genBoundsOps(builder, varPtr, boundsOps);
 
+    mlir::omp::VariableCaptureKind captureKind =
+        mlir::omp::VariableCaptureKind::ByRef;
+    if (fir::isa_trivial(fir::unwrapRefType(varPtr.getType())) ||
+        fir::isa_char(fir::unwrapRefType(varPtr.getType()))) {
+      if (isLiteralType(fir::unwrapRefType(varPtr.getType()))) {
+        captureKind = mlir::omp::VariableCaptureKind::ByCopy;
+      }
+    }
+
     return builder.create<omp::MapInfoOp>(
         loc, varPtr.getType(), varPtr,
         TypeAttr::get(llvm::cast<omp::PointerLikeType>(varPtr.getType())
                           .getElementType()),
         builder.getIntegerAttr(builder.getIntegerType(64, /*isSigned=*/false),
                                mapTypeTo),
-        builder.getAttr<omp::VariableCaptureKindAttr>(
-            omp::VariableCaptureKind::ByRef),
+        builder.getAttr<omp::VariableCaptureKindAttr>(captureKind),
         /*varPtrPtr=*/Value{},
         /*members=*/SmallVector<Value>{},
         /*member_index=*/mlir::ArrayAttr{},