[flang][OpenMP] Move ALLOCATE + privatize check to semantic checks (#192792)

Move the check from symbol resolution to semantic checks.

The check now seems to be more accurate, catching some cases that were
not detected before.

Author: kparzysz

flang/include/flang/Parser/openmp-utils.h

@@ -36,6 +36,34 @@ template <typename T> constexpr auto addr_if(const std::optional<T> &x) {
   return x ? &*x : nullptr;
 }
 
+namespace detail {
+struct DirectiveSpecificationScope {
+  using ODS = OmpDirectiveSpecification;
+  template <typename T> static const ODS &GetODS(const T &x) {
+    if constexpr ( //
+        std::is_base_of_v<OmpBlockConstruct, T> ||
+        std::is_same_v<OpenMPLoopConstruct, T> ||
+        std::is_same_v<OpenMPSectionsConstruct, T>) {
+      return x.BeginDir();
+    } else if constexpr (WrapperTrait<T>) {
+      return GetODS(x.v);
+    } else if constexpr (UnionTrait<T>) {
+      return std::visit(
+          [](auto &&s) -> decltype(auto) { return GetODS(s); }, x.u);
+    } else {
+      static_assert(std::is_same_v<OpenMPSectionConstruct, T>);
+      llvm_unreachable("This function does not work for SECTION");
+    }
+  }
+  static inline const ODS &GetODS(const ODS &x) { return x; }
+};
+} // namespace detail
+
+const OmpDirectiveSpecification &GetOmpDirectiveSpecification(
+    const OpenMPConstruct &x);
+const OmpDirectiveSpecification &GetOmpDirectiveSpecification(
+    const OpenMPDeclarativeConstruct &x);
+
 namespace detail {
 struct DirectiveNameScope {
   static OmpDirectiveName MakeName(CharBlock source = {},

flang/lib/Parser/openmp-utils.cpp

@@ -25,6 +25,24 @@
 
 namespace Fortran::parser::omp {
 
+const OmpDirectiveSpecification &GetOmpDirectiveSpecification(
+    const OpenMPConstruct &x) {
+  return std::visit(
+      [](auto &&s) -> decltype(auto) {
+        return detail::DirectiveSpecificationScope::GetODS(s);
+      },
+      x.u);
+}
+
+const OmpDirectiveSpecification &GetOmpDirectiveSpecification(
+    const OpenMPDeclarativeConstruct &x) {
+  return std::visit(
+      [](auto &&s) -> decltype(auto) {
+        return detail::DirectiveSpecificationScope::GetODS(s);
+      },
+      x.u);
+}
+
 std::string GetUpperName(llvm::omp::Clause id, unsigned version) {
   llvm::StringRef name{llvm::omp::getOpenMPClauseName(id, version)};
   return parser::ToUpperCaseLetters(name);

flang/lib/Semantics/check-omp-structure.cpp

@@ -336,6 +336,16 @@ class OmpUnitedTaskDesignatorChecker {
   SemanticsContext &context_;
 };
 
+bool OmpStructureChecker::IsAllowedClause(llvm::omp::Clause clauseId) {
+  // Do not do clause checks while processing METADIRECTIVE.
+  // See comment in CheckAllowedClause.
+  if (GetDirectiveNest(ContextSelectorNest) > 0) {
+    return true;
+  }
+  return llvm::omp::isAllowedClauseForDirective(
+      GetContext().directive, clauseId, context_.langOptions().OpenMPVersion);
+}
+
 bool OmpStructureChecker::CheckAllowedClause(llvmOmpClause clause) {
   // Do not do clause checks while processing METADIRECTIVE.
   // Context selectors can contain clauses that are not given as a part
@@ -772,6 +782,9 @@ void OmpStructureChecker::Enter(const parser::OpenMPConstruct &x) {
         return CheckDirectiveSpelling(source, id);
       });
   parser::Walk(x, visitor);
+  if (GetOmpDirectiveName(x).v != llvm::omp::Directive::OMPD_section) {
+    dirStack_.push_back(&GetOmpDirectiveSpecification(x));
+  }
 
   // Simd Construct with Ordered Construct Nesting check
   // We cannot use CurrentDirectiveIsNested() here because
@@ -788,12 +801,15 @@ void OmpStructureChecker::Enter(const parser::OpenMPConstruct &x) {
   }
 }
 
-void OmpStructureChecker::Leave(const parser::OpenMPConstruct &) {
+void OmpStructureChecker::Leave(const parser::OpenMPConstruct &x) {
   for (const auto &[sym, source] : deferredNonVariables_) {
     context_.SayWithDecl(
         *sym, source, "'%s' must be a variable"_err_en_US, sym->name());
   }
   deferredNonVariables_.clear();
+  if (GetOmpDirectiveName(x).v != llvm::omp::Directive::OMPD_section) {
+    dirStack_.pop_back();
+  }
 }
 
 void OmpStructureChecker::Enter(const parser::OpenMPDeclarativeConstruct &x) {
@@ -803,11 +819,13 @@ void OmpStructureChecker::Enter(const parser::OpenMPDeclarativeConstruct &x) {
       });
   parser::Walk(x, visitor);
 
+  dirStack_.push_back(&GetOmpDirectiveSpecification(x));
   EnterDirectiveNest(DeclarativeNest);
 }
 
 void OmpStructureChecker::Leave(const parser::OpenMPDeclarativeConstruct &x) {
   ExitDirectiveNest(DeclarativeNest);
+  dirStack_.pop_back();
 }
 
 void OmpStructureChecker::AddEndDirectiveClauses(
@@ -2010,6 +2028,50 @@ void OmpStructureChecker::Enter(const parser::OmpClause::Allocate &x) {
       }
     }
   }
+
+  // If the clause is not allowed, don't diagnose specific problems with it.
+  if (!IsAllowedClause(llvm::omp::Clause::OMPC_allocate)) {
+    return;
+  }
+
+  llvm::omp::Directive dirId{GetContext().directive};
+  // For any list item that is specified in the allocate clause on a directive
+  // other than the allocators directive, a data-sharing attribute clause that
+  // may create a private copy of that list item must be specified on the same
+  // directive.
+  if (dirId != llvm::omp::Directive::OMPD_allocators &&
+      dirId != llvm::omp::Directive::OMPD_section) {
+    static llvm::omp::Clause privatizingDSAClauses[] = {
+        llvm::omp::Clause::OMPC_private,
+        llvm::omp::Clause::OMPC_firstprivate,
+        llvm::omp::Clause::OMPC_lastprivate,
+        llvm::omp::Clause::OMPC_is_device_ptr,
+        llvm::omp::Clause::OMPC_use_device_ptr,
+    };
+    const parser::OmpDirectiveSpecification &spec{*dirStack_.back()};
+
+    std::set<const Symbol *> privatized;
+    for (auto dsaClause : privatizingDSAClauses) {
+      if (auto *found{parser::omp::FindClause(spec, dsaClause)}) {
+        for (auto &object : GetOmpObjectList(*found)->v) {
+          if (auto *symbol{GetObjectSymbol(object)}) {
+            privatized.insert(symbol);
+          }
+        }
+      }
+    }
+
+    for (auto &object : GetOmpObjectList(x)->v) {
+      if (auto *symbol{GetObjectSymbol(object)}) {
+        if (!privatized.count(symbol)) {
+          context_.Say(
+              GetObjectSource(object).value_or(GetContext().clauseSource),
+              "The ALLOCATE clause requires that '%s' must be listed in a private data-sharing attribute clause on the same directive"_err_en_US,
+              symbol->name());
+        }
+      }
+    }
+  }
 }
 
 void OmpStructureChecker::Enter(const parser::OpenMPDeclareMapperConstruct &x) {

flang/lib/Semantics/check-omp-structure.h

@@ -93,8 +93,6 @@ class OmpStructureChecker : public OmpStructureCheckerBase {
 
   void Enter(const parser::OpenMPConstruct &);
   void Leave(const parser::OpenMPConstruct &);
-  void Enter(const parser::OpenMPInteropConstruct &);
-  void Leave(const parser::OpenMPInteropConstruct &);
   void Enter(const parser::OpenMPDeclarativeConstruct &);
   void Leave(const parser::OpenMPDeclarativeConstruct &);
 
@@ -112,6 +110,8 @@ class OmpStructureChecker : public OmpStructureCheckerBase {
   void Leave(const parser::OpenMPAssumeConstruct &);
   void Enter(const parser::OpenMPDeclarativeAssumes &);
   void Leave(const parser::OpenMPDeclarativeAssumes &);
+  void Enter(const parser::OpenMPInteropConstruct &);
+  void Leave(const parser::OpenMPInteropConstruct &);
   void Enter(const parser::OmpBlockConstruct &);
   void Leave(const parser::OmpBlockConstruct &);
   void Leave(const parser::OmpBeginDirective &);
@@ -267,6 +267,7 @@ class OmpStructureChecker : public OmpStructureCheckerBase {
       const parser::OmpTraitSetSelector &, const parser::OmpTraitSelector &);
 
   // check-omp-structure.cpp
+  bool IsAllowedClause(llvm::omp::Clause clauseId);
   bool CheckAllowedClause(llvmOmpClause clause);
   void CheckVariableListItem(const SymbolSourceMap &symbols);
   void CheckDirectiveSpelling(
@@ -404,6 +405,10 @@ class OmpStructureChecker : public OmpStructureCheckerBase {
   std::vector<LoopConstruct> loopStack_;
   // Scopes for scoping units.
   std::vector<const Scope *> scopeStack_;
+  // Stack of directive specifications (except for SECTION).
+  // This is to allow visitor functions to see all specified clauses, since
+  // they are only recorded in DirContext as they are processed.
+  std::vector<const parser::OmpDirectiveSpecification *> dirStack_;
 
   enum class PartKind : int {
     // There are also other "parts", such as internal-subprogram-part, etc,

flang/lib/Semantics/resolve-directives.cpp

@@ -1082,7 +1082,6 @@ class OmpAttributeVisitor : DirectiveAttributeVisitor<llvm::omp::Directive> {
   Symbol::Flags dataCopyingAttributeFlags{
       Symbol::Flag::OmpCopyIn, Symbol::Flag::OmpCopyPrivate};
 
-  std::vector<const parser::Name *> allocateNames_; // on one directive
   UnorderedSymbolSet privateDataSharingAttributeObjects_; // on one directive
   UnorderedSymbolSet stmtFunctionExprSymbols_;
   std::multimap<const parser::Label,
@@ -1101,11 +1100,6 @@ class OmpAttributeVisitor : DirectiveAttributeVisitor<llvm::omp::Directive> {
   };
   std::vector<PartKind> partStack_;
 
-  void AddAllocateName(const parser::Name *&object) {
-    allocateNames_.push_back(object);
-  }
-  void ClearAllocateNames() { allocateNames_.clear(); }
-
   void AddPrivateDataSharingAttributeObjects(SymbolRef object) {
     privateDataSharingAttributeObjects_.insert(object);
   }
@@ -2006,35 +2000,10 @@ bool OmpAttributeVisitor::Pre(const parser::OmpBlockConstruct &x) {
     IssueNonConformanceWarning(dirId, dirSpec.source, 52);
   ClearDataSharingAttributeObjects();
   ClearPrivateDataSharingAttributeObjects();
-  ClearAllocateNames();
   return true;
 }
 
 void OmpAttributeVisitor::Post(const parser::OmpBlockConstruct &x) {
-  const parser::OmpDirectiveSpecification &dirSpec{x.BeginDir()};
-  llvm::omp::Directive dirId{dirSpec.DirId()};
-  unsigned version{context_.langOptions().OpenMPVersion};
-
-  if (llvm::omp::isPrivatizingConstruct(dirId, version)) {
-    bool hasPrivate;
-    for (const auto *allocName : allocateNames_) {
-      hasPrivate = false;
-      for (auto privateObj : privateDataSharingAttributeObjects_) {
-        const Symbol &symbolPrivate{*privateObj};
-        if (allocName->source == symbolPrivate.name()) {
-          hasPrivate = true;
-          break;
-        }
-      }
-      if (!hasPrivate) {
-        context_.Say(allocName->source,
-            "The ALLOCATE clause requires that '%s' must be listed in a "
-            "private "
-            "data-sharing attribute clause on the same directive"_err_en_US,
-            allocName->ToString());
-      }
-    }
-  }
   PopContext();
 }
 
@@ -2931,10 +2900,6 @@ void OmpAttributeVisitor::ResolveOmpDesignator(
       if (privateDataSharingAttributeFlags.test(ompFlag)) {
         CheckObjectIsPrivatizable(*name, *symbol, ompFlag);
       }
-
-      if (ompFlag == Symbol::Flag::OmpAllocate) {
-        AddAllocateName(name);
-      }
     }
     // Save the original symbol. For privatizing clauses, ensure enclosing
     // constructs properly capture the variable.

flang/test/Lower/OpenMP/sections.f90

@@ -12,13 +12,15 @@
 !CHECK:   %[[CONST_1:.*]] = arith.constant 4 : i64
 !CHECK:   %[[PRIVATE_ETA:.*]] = fir.alloca f32 {bindc_name = "eta", pinned, uniq_name = "_QFEeta"}
 !CHECK:   %[[PRIVATE_ETA_DECL:.*]]:2 = hlfir.declare %[[PRIVATE_ETA]] {uniq_name = "_QFEeta"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
+!CHECK:   %[[PRIVATE_COUNT:.*]] = fir.alloca i32 {bindc_name = "count", pinned, uniq_name = "_QFEcount"}
+!CHECK:   %[[PRIVATE_COUNT_DECL:.*]]:2 = hlfir.declare %[[PRIVATE_COUNT]] {uniq_name = "_QFEcount"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
 !CHECK:   %[[PRIVATE_DOUBLE_COUNT:.*]] = fir.alloca i32 {bindc_name = "double_count", pinned, uniq_name = "_QFEdouble_count"}
 !CHECK:   %[[PRIVATE_DOUBLE_COUNT_DECL:.*]]:2 = hlfir.declare %[[PRIVATE_DOUBLE_COUNT]] {uniq_name = "_QFEdouble_count"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
 !CHECK:   omp.sections allocate(%[[CONST_1]] : i64 -> %[[COUNT_DECL]]#0 : !fir.ref<i32>)  {
 !CHECK:     omp.section {
 !CHECK:       %[[CONST5:.*]] = arith.constant 5 : i32
-!CHECK:       hlfir.assign %[[CONST5]] to %[[COUNT_DECL]]#0 : i32, !fir.ref<i32>
-!CHECK:       %[[TEMP_COUNT:.*]] = fir.load %[[COUNT_DECL]]#0 : !fir.ref<i32>
+!CHECK:       hlfir.assign %[[CONST5]] to %[[PRIVATE_COUNT_DECL]]#0 : i32, !fir.ref<i32>
+!CHECK:       %[[TEMP_COUNT:.*]] = fir.load %[[PRIVATE_COUNT_DECL]]#0 : !fir.ref<i32>
 !CHECK:       %[[TEMP_DOUBLE_COUNT:.*]] = fir.load %[[PRIVATE_DOUBLE_COUNT_DECL]]#0 : !fir.ref<i32>
 !CHECK:       %[[RESULT:.*]] = arith.muli %[[TEMP_COUNT]], %[[TEMP_DOUBLE_COUNT]] : i32
 !CHECK:       %[[RESULT_CONVERT:.*]] = fir.convert %[[RESULT]] : (i32) -> f32
@@ -37,13 +39,13 @@
 !CHECK:       %[[CONST:.*]] = arith.constant 7.000000e+00 : f32
 !CHECK:       %[[RESULT:.*]] = arith.subf %[[TEMP]], %[[CONST]] {{.*}}: f32
 !CHECK:       hlfir.assign %[[RESULT]] to %[[PRIVATE_ETA_DECL]]#0 : f32, !fir.ref<f32>
-!CHECK:       %[[TEMP_COUNT1:.*]] = fir.load %[[COUNT_DECL]]#0 : !fir.ref<i32>
+!CHECK:       %[[TEMP_COUNT1:.*]] = fir.load %[[PRIVATE_COUNT_DECL]]#0 : !fir.ref<i32>
 !CHECK:       %[[TEMP_COUNT:.*]] = fir.convert %[[TEMP_COUNT1]] : (i32) -> f32
 !CHECK:       %[[TEMP_ETA:.*]] = fir.load %[[PRIVATE_ETA_DECL]]#0 : !fir.ref<f32>
 !CHECK:       %[[TEMP_COUNT2:.*]] = arith.mulf %[[TEMP_COUNT]], %[[TEMP_ETA]] {{.*}}: f32
 !CHECK:       %[[RESULT:.*]] = fir.convert %[[TEMP_COUNT2]] : (f32) -> i32
-!CHECK:       hlfir.assign %[[RESULT]] to %[[COUNT_DECL]]#0 : i32, !fir.ref<i32>
-!CHECK:       %[[TEMP_COUNT3:.*]] = fir.load %[[COUNT_DECL]]#0 : !fir.ref<i32>
+!CHECK:       hlfir.assign %[[RESULT]] to %[[PRIVATE_COUNT_DECL]]#0 : i32, !fir.ref<i32>
+!CHECK:       %[[TEMP_COUNT3:.*]] = fir.load %[[PRIVATE_COUNT_DECL]]#0 : !fir.ref<i32>
 !CHECK:       %[[TEMP_COUNT4:.*]] = fir.convert %[[TEMP_COUNT3]] : (i32) -> f32
 !CHECK:       %[[TEMP_ETA:.*]] = fir.load %[[PRIVATE_ETA_DECL]]#0 : !fir.ref<f32>
 !CHECK:       %[[TEMP_COUNT5:.*]] = arith.subf %[[TEMP_COUNT4]], %[[TEMP_ETA]] {{.*}}: f32
@@ -62,7 +64,7 @@
 program sample
     use omp_lib
     integer :: count = 0, double_count = 1
-    !$omp sections private (eta, double_count) allocate(omp_high_bw_mem_alloc: count)
+    !$omp sections private (eta, count, double_count) allocate(omp_high_bw_mem_alloc: count)
         !$omp section
             count = 1 + 4
             eta = count * double_count

flang/test/Semantics/OpenMP/taskgroup01.f90

@@ -21,9 +21,11 @@
   !$omp end parallel
 
   !$omp parallel private(xyz)
+  !ERROR: The ALLOCATE clause requires that 'xyz' must be listed in a private data-sharing attribute clause on the same directive
     !$omp taskgroup allocate(xyz)
       !$omp task
         print *, "The "
+  !ERROR: The ALLOCATE clause requires that 'abc' must be listed in a private data-sharing attribute clause on the same directive
         !$omp taskgroup allocate(omp_large_cap_mem_space: abc)
           !$omp task
           print *, "almighty sun"