[flang][OpenMP] Generalize isOpenMPPrivatizingConstruct (#148654)

Instead of treating all block and all loop constructs as privatizing,
actually check if the construct allows a privatizing clause.

Author: kparzysz

flang/lib/Lower/OpenMP/DataSharingProcessor.cpp

@@ -26,6 +26,8 @@
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Semantics/attr.h"
 #include "flang/Semantics/tools.h"
+#include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/SmallSet.h"
 
 namespace Fortran {
 namespace lower {
@@ -49,7 +51,7 @@ DataSharingProcessor::DataSharingProcessor(
       firOpBuilder(converter.getFirOpBuilder()), clauses(clauses), eval(eval),
       shouldCollectPreDeterminedSymbols(shouldCollectPreDeterminedSymbols),
       useDelayedPrivatization(useDelayedPrivatization), symTable(symTable),
-      visitor() {
+      visitor(semaCtx) {
   eval.visit([&](const auto &functionParserNode) {
     parser::Walk(functionParserNode, visitor);
   });
@@ -424,24 +426,55 @@ getSource(const semantics::SemanticsContext &semaCtx,
   return source;
 }
 
+static void collectPrivatizingConstructs(
+    llvm::SmallSet<llvm::omp::Directive, 16> &constructs, unsigned version) {
+  using Clause = llvm::omp::Clause;
+  using Directive = llvm::omp::Directive;
+
+  static const Clause privatizingClauses[] = {
+      Clause::OMPC_private,
+      Clause::OMPC_lastprivate,
+      Clause::OMPC_firstprivate,
+      Clause::OMPC_in_reduction,
+      Clause::OMPC_reduction,
+      Clause::OMPC_linear,
+      // TODO: Clause::OMPC_induction,
+      Clause::OMPC_task_reduction,
+      Clause::OMPC_detach,
+      Clause::OMPC_use_device_ptr,
+      Clause::OMPC_is_device_ptr,
+  };
+
+  for (auto dir : llvm::enum_seq_inclusive<Directive>(Directive::First_,
+                                                      Directive::Last_)) {
+    bool allowsPrivatizing = llvm::any_of(privatizingClauses, [&](Clause cls) {
+      return llvm::omp::isAllowedClauseForDirective(dir, cls, version);
+    });
+    if (allowsPrivatizing)
+      constructs.insert(dir);
+  }
+}
+
 bool DataSharingProcessor::isOpenMPPrivatizingConstruct(
-    const parser::OpenMPConstruct &omp) {
-  return common::visit(
-      [](auto &&s) {
-        using BareS = llvm::remove_cvref_t<decltype(s)>;
-        return std::is_same_v<BareS, parser::OpenMPBlockConstruct> ||
-               std::is_same_v<BareS, parser::OpenMPLoopConstruct> ||
-               std::is_same_v<BareS, parser::OpenMPSectionsConstruct>;
-      },
-      omp.u);
+    const parser::OpenMPConstruct &omp, unsigned version) {
+  static llvm::SmallSet<llvm::omp::Directive, 16> privatizing;
+  [[maybe_unused]] static bool init =
+      (collectPrivatizingConstructs(privatizing, version), true);
+
+  // As of OpenMP 6.0, privatizing constructs (with the test being if they
+  // allow a privatizing clause) are: dispatch, distribute, do, for, loop,
+  // parallel, scope, sections, simd, single, target, target_data, task,
+  // taskgroup, taskloop, and teams.
+  return llvm::is_contained(privatizing, extractOmpDirective(omp));
 }
 
 bool DataSharingProcessor::isOpenMPPrivatizingEvaluation(
     const pft::Evaluation &eval) const {
-  return eval.visit([](auto &&s) {
+  unsigned version = semaCtx.langOptions().OpenMPVersion;
+  return eval.visit([=](auto &&s) {
     using BareS = llvm::remove_cvref_t<decltype(s)>;
     if constexpr (std::is_same_v<BareS, parser::OpenMPConstruct>) {
-      return isOpenMPPrivatizingConstruct(s);
+      return isOpenMPPrivatizingConstruct(s, version);
     } else {
       return false;
     }

flang/lib/Lower/OpenMP/DataSharingProcessor.h

@@ -36,6 +36,8 @@ class DataSharingProcessor {
   /// at any point in time. This is used to track Symbol definition scopes in
   /// order to tell which OMP scope defined vs. references a certain Symbol.
   struct OMPConstructSymbolVisitor {
+    OMPConstructSymbolVisitor(semantics::SemanticsContext &ctx)
+        : version(ctx.langOptions().OpenMPVersion) {}
     template <typename T>
     bool Pre(const T &) {
       return true;
@@ -45,13 +47,13 @@ class DataSharingProcessor {
 
     bool Pre(const parser::OpenMPConstruct &omp) {
       // Skip constructs that may not have privatizations.
-      if (isOpenMPPrivatizingConstruct(omp))
+      if (isOpenMPPrivatizingConstruct(omp, version))
         constructs.push_back(&omp);
       return true;
     }
 
     void Post(const parser::OpenMPConstruct &omp) {
-      if (isOpenMPPrivatizingConstruct(omp))
+      if (isOpenMPPrivatizingConstruct(omp, version))
         constructs.pop_back();
     }
 
@@ -68,6 +70,9 @@ class DataSharingProcessor {
     /// construct that defines symbol.
     bool isSymbolDefineBy(const semantics::Symbol *symbol,
                           lower::pft::Evaluation &eval) const;
+
+  private:
+    unsigned version;
   };
 
   mlir::OpBuilder::InsertPoint lastPrivIP;
@@ -115,7 +120,8 @@ class DataSharingProcessor {
                              mlir::OpBuilder::InsertPoint *lastPrivIP);
   void insertDeallocs();
 
-  static bool isOpenMPPrivatizingConstruct(const parser::OpenMPConstruct &omp);
+  static bool isOpenMPPrivatizingConstruct(const parser::OpenMPConstruct &omp,
+                                           unsigned version);
   bool isOpenMPPrivatizingEvaluation(const pft::Evaluation &eval) const;
 
 public:

flang/test/Lower/OpenMP/taskgroup02.f90

@@ -3,8 +3,9 @@
 ! Check that variables are not privatized twice when TASKGROUP is used.
 
 !CHECK-LABEL: func.func @_QPsub() {
-!CHECK:         omp.parallel {
-!CHECK:           %[[PAR_I:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "_QFsubEi"}
+!CHECK:         omp.parallel private(@_QFsubEi_private_i32 %[[SUB_I:.*]]#0 -> %[[ARG:.*]] : !fir.ref<i32>)
+!CHECK:           %[[ALLOCA:.*]] = fir.alloca i32
+!CHECK:           %[[PAR_I:.*]]:2 = hlfir.declare %[[ALLOCA]] {uniq_name = "_QFsubEi"}
 !CHECK:           omp.master {
 !CHECK:             omp.taskgroup {
 !CHECK-NEXT:          omp.task private(@_QFsubEi_firstprivate_i32 %[[PAR_I]]#0 -> %[[TASK_I:.*]] : !fir.ref<i32>) {

llvm/include/llvm/Frontend/OpenMP/OMP.h

@@ -51,13 +51,17 @@ static constexpr inline bool canHaveIterator(Clause C) {
 // Can clause C create a private copy of a variable.
 static constexpr inline bool isPrivatizingClause(Clause C) {
   switch (C) {
+  case OMPC_detach:
   case OMPC_firstprivate:
+  // TODO case OMPC_induction:
   case OMPC_in_reduction:
+  case OMPC_is_device_ptr:
   case OMPC_lastprivate:
   case OMPC_linear:
   case OMPC_private:
   case OMPC_reduction:
   case OMPC_task_reduction:
+  case OMPC_use_device_ptr:
     return true;
   default:
     return false;