[flang][OpenMP] Try to find a proper solution for symbol scoping in DSP

Author: ergawy

flang/lib/Lower/OpenMP/DataSharingProcessor.cpp

@@ -22,6 +22,30 @@
 namespace Fortran {
 namespace lower {
 namespace omp {
+bool DataSharingProcessor::OMPConstructSymbolVisitor::isSymbolDefineBy(
+    const semantics::Symbol *symbol, lower::pft::Evaluation &eval) const {
+  return eval.visit(
+      common::visitors{[&](const parser::OpenMPConstruct &functionParserNode) {
+                         return symDefMap.count(symbol) &&
+                                symDefMap.at(symbol) == &functionParserNode;
+                       },
+                       [](const auto &functionParserNode) { 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)
+    : hasLastPrivateOp(false), converter(converter), semaCtx(semaCtx),
+      firOpBuilder(converter.getFirOpBuilder()), clauses(clauses), eval(eval),
+      shouldCollectPreDeterminedSymbols(shouldCollectPreDeterminedSymbols),
+      useDelayedPrivatization(useDelayedPrivatization), symTable(symTable),
+      visitor() {
+  eval.visit([&](const auto &functionParserNode) {
+    parser::Walk(functionParserNode, visitor);
+  });
+}
 
 void DataSharingProcessor::processStep1(
     mlir::omp::PrivateClauseOps *clauseOps,
@@ -285,38 +309,9 @@ void DataSharingProcessor::collectSymbolsInNestedRegions(
         // Recursively look for OpenMP constructs within `nestedEval`'s region
         collectSymbolsInNestedRegions(nestedEval, flag, symbolsInNestedRegions);
       else {
-        bool isOrderedConstruct = [&]() {
-          if (auto *ompConstruct =
-                  nestedEval.getIf<parser::OpenMPConstruct>()) {
-            if (auto *ompBlockConstruct =
-                    std::get_if<parser::OpenMPBlockConstruct>(
-                        &ompConstruct->u)) {
-              const auto &beginBlockDirective =
-                  std::get<parser::OmpBeginBlockDirective>(
-                      ompBlockConstruct->t);
-              const auto origDirective =
-                  std::get<parser::OmpBlockDirective>(beginBlockDirective.t).v;
-
-              return origDirective == llvm::omp::Directive::OMPD_ordered;
-            }
-          }
-
-          return false;
-        }();
-
-        bool isCriticalConstruct = [&]() {
-          if (auto *ompConstruct =
-                  nestedEval.getIf<parser::OpenMPConstruct>()) {
-            return std::get_if<parser::OpenMPCriticalConstruct>(
-                       &ompConstruct->u) != nullptr;
-          }
-          return false;
-        }();
-
-        if (!isOrderedConstruct && !isCriticalConstruct)
-          converter.collectSymbolSet(nestedEval, symbolsInNestedRegions, flag,
-                                     /*collectSymbols=*/true,
-                                     /*collectHostAssociatedSymbols=*/false);
+        converter.collectSymbolSet(nestedEval, symbolsInNestedRegions, flag,
+                                   /*collectSymbols=*/true,
+                                   /*collectHostAssociatedSymbols=*/false);
       }
     }
   }
@@ -356,6 +351,11 @@ void DataSharingProcessor::collectSymbols(
 
   llvm::SetVector<const semantics::Symbol *> symbolsInNestedRegions;
   collectSymbolsInNestedRegions(eval, flag, symbolsInNestedRegions);
+
+  for (auto *symbol : allSymbols)
+    if (visitor.isSymbolDefineBy(symbol, eval))
+      symbolsInNestedRegions.remove(symbol);
+
   // Filter-out symbols that must not be privatized.
   bool collectImplicit = flag == semantics::Symbol::Flag::OmpImplicit;
   bool collectPreDetermined = flag == semantics::Symbol::Flag::OmpPreDetermined;

flang/lib/Lower/OpenMP/DataSharingProcessor.h

@@ -32,6 +32,40 @@ namespace omp {
 
 class DataSharingProcessor {
 private:
+  /// A symbol visitor that keeps track of the currently active OpenMPConstruct
+  /// 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 {
+    template <typename T>
+    bool Pre(const T &) {
+      return true;
+    }
+    template <typename T>
+    void Post(const T &) {}
+
+    bool Pre(const parser::OpenMPConstruct &omp) {
+      currentConstruct = &omp;
+      return true;
+    }
+
+    void Post(const parser::OpenMPConstruct &omp) {
+      currentConstruct = nullptr;
+    }
+
+    void Post(const parser::Name &name) {
+      symDefMap.try_emplace(name.symbol, currentConstruct);
+    }
+
+    const parser::OpenMPConstruct *currentConstruct = nullptr;
+    llvm::DenseMap<semantics::Symbol *, const parser::OpenMPConstruct *>
+        symDefMap;
+
+    /// Given a \p symbol and an \p eval, returns true if eval is the OMP
+    /// construct that defines symbol.
+    bool isSymbolDefineBy(const semantics::Symbol *symbol,
+                          lower::pft::Evaluation &eval) const;
+  };
+
   bool hasLastPrivateOp;
   mlir::OpBuilder::InsertPoint lastPrivIP;
   mlir::OpBuilder::InsertPoint insPt;
@@ -53,6 +87,7 @@ class DataSharingProcessor {
   bool shouldCollectPreDeterminedSymbols;
   bool useDelayedPrivatization;
   lower::SymMap *symTable;
+  OMPConstructSymbolVisitor visitor;
 
   bool needBarrier();
   void collectSymbols(semantics::Symbol::Flag flag,
@@ -97,11 +132,7 @@ class DataSharingProcessor {
                        lower::pft::Evaluation &eval,
                        bool shouldCollectPreDeterminedSymbols,
                        bool useDelayedPrivatization = false,
-                       lower::SymMap *symTable = nullptr)
-      : hasLastPrivateOp(false), converter(converter), semaCtx(semaCtx),
-        firOpBuilder(converter.getFirOpBuilder()), clauses(clauses), eval(eval),
-        shouldCollectPreDeterminedSymbols(shouldCollectPreDeterminedSymbols),
-        useDelayedPrivatization(useDelayedPrivatization), symTable(symTable) {}
+                       lower::SymMap *symTable = nullptr);
 
   // Privatisation is split into two steps.
   // Step1 performs cloning of all privatisation clauses and copying for