[flang][OpenMP] Fix detecting nested OpenMP constructs

flang/lib/Lower/OpenMP/DataSharingProcessor.cpp

@@ -388,19 +388,43 @@ getSource(const semantics::SemanticsContext &semaCtx,
   return source;
 }
 
+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);
+}
+
+bool DataSharingProcessor::isOpenMPPrivatizingEvaluation(
+    const pft::Evaluation &eval) const {
+  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);
+    } else {
+      return false;
+    }
+  });
+}
+
 void DataSharingProcessor::collectSymbolsInNestedRegions(
     lower::pft::Evaluation &eval, semantics::Symbol::Flag flag,
     llvm::SetVector<const semantics::Symbol *> &symbolsInNestedRegions) {
-  for (lower::pft::Evaluation &nestedEval : eval.getNestedEvaluations()) {
-    if (nestedEval.hasNestedEvaluations()) {
-      if (nestedEval.isConstruct())
-        // Recursively look for OpenMP constructs within `nestedEval`'s region
-        collectSymbolsInNestedRegions(nestedEval, flag, symbolsInNestedRegions);
-      else {
-        converter.collectSymbolSet(nestedEval, symbolsInNestedRegions, flag,
-                                   /*collectSymbols=*/true,
-                                   /*collectHostAssociatedSymbols=*/false);
-      }
+  if (!eval.hasNestedEvaluations())
+    return;
+  for (pft::Evaluation &nestedEval : eval.getNestedEvaluations()) {
+    if (isOpenMPPrivatizingEvaluation(nestedEval)) {
+      converter.collectSymbolSet(nestedEval, symbolsInNestedRegions, flag,
+                                 /*collectSymbols=*/true,
+                                 /*collectHostAssociatedSymbols=*/false);
+    } else {
+      // Recursively look for OpenMP constructs within `nestedEval`'s region
+      collectSymbolsInNestedRegions(nestedEval, flag, symbolsInNestedRegions);
     }
   }
 }

flang/lib/Lower/OpenMP/DataSharingProcessor.h

@@ -45,20 +45,22 @@ class DataSharingProcessor {
 
     bool Pre(const parser::OpenMPConstruct &omp) {
       // Skip constructs that may not have privatizations.
-      if (!std::holds_alternative<parser::OpenMPCriticalConstruct>(omp.u))
-        currentConstruct = &omp;
+      if (isOpenMPPrivatizingConstruct(omp))
+        constructs.push_back(&omp);
       return true;
     }
 
     void Post(const parser::OpenMPConstruct &omp) {
-      currentConstruct = nullptr;
+      if (isOpenMPPrivatizingConstruct(omp))
+        constructs.pop_back();
     }
 
     void Post(const parser::Name &name) {
-      symDefMap.try_emplace(name.symbol, currentConstruct);
+      auto *current = !constructs.empty() ? constructs.back() : nullptr;
+      symDefMap.try_emplace(name.symbol, current);
     }
 
-    const parser::OpenMPConstruct *currentConstruct = nullptr;
+    llvm::SmallVector<const parser::OpenMPConstruct *> constructs;
     llvm::DenseMap<semantics::Symbol *, const parser::OpenMPConstruct *>
         symDefMap;
 
@@ -113,6 +115,9 @@ class DataSharingProcessor {
                              mlir::OpBuilder::InsertPoint *lastPrivIP);
   void insertDeallocs();
 
+  static bool isOpenMPPrivatizingConstruct(const parser::OpenMPConstruct &omp);
+  bool isOpenMPPrivatizingEvaluation(const pft::Evaluation &eval) const;
+
 public:
   DataSharingProcessor(lower::AbstractConverter &converter,
                        semantics::SemanticsContext &semaCtx,

flang/test/Lower/OpenMP/atomic-privatize.f90

@@ -0,0 +1,25 @@
+! Testcase adopted from the Fujitsu test suite:
+! https://github.com/fujitsu/compiler-test-suite/blob/main/Fortran/0407/0407_0006.f90
+
+!RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=60 %s -o - | FileCheck %s
+
+! Make sure that the variable in the atomic construct is privatized at
+! the task level
+
+!CHECK: omp.task private(@_QFfredEprv_firstprivate_i32 %{{[0-9]+}}#0 -> %arg0
+!CHECK: %[[DECL:[0-9]+]]:2 = hlfir.declare %arg0 {uniq_name = "_QFfredEprv"}
+!CHECK: omp.atomic.update %[[DECL]]#0
+
+integer function fred
+  integer :: prv
+
+  prv = 1
+  !$omp parallel shared(prv)
+  !$omp task default(firstprivate)
+  !$omp atomic
+  prv = prv + 1
+  !$omp end task
+  !$omp end parallel
+  fred = prv
+end
+

flang/test/Lower/OpenMP/sections-predetermined-private.f90

@@ -11,7 +11,7 @@
 end
 ! CHECK-LABEL:   func.func @_QQmain() {
 ! CHECK:           omp.parallel {
-! CHECK:             %[[VAL_3:.*]] = fir.alloca i32 {bindc_name = "i", pinned}
+! CHECK:             %[[VAL_3:.*]] = fir.alloca i32 {bindc_name = "i", pinned, uniq_name = "_QFEi"}
 ! CHECK:             %[[VAL_4:.*]]:2 = hlfir.declare %[[VAL_3]] {uniq_name = "_QFEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
 ! CHECK:             omp.sections {
 ! CHECK:               omp.section {