[flang][openacc] Fix semantic check for wait and atomic directives

wait and atomic directives are represented by OpenACCWaitConstruct, OpenACCAtmicConstruct in the parser. Those contrsuct were
not taken into account in the semantic check so far.

Reviewed By: kiranchandramohan

Differential Revision: https://reviews.llvm.org/D88628

flang/include/flang/Parser/parse-tree.h

@@ -3970,6 +3970,7 @@ struct OpenACCAtomicConstruct {
   UNION_CLASS_BOILERPLATE(OpenACCAtomicConstruct);
   std::variant<AccAtomicRead, AccAtomicWrite, AccAtomicCapture, AccAtomicUpdate>
       u;
+  CharBlock source;
 };
 
 struct OpenACCBlockConstruct {

flang/lib/Parser/openacc-parsers.cpp

@@ -242,10 +242,11 @@ TYPE_PARSER("ATOMIC" >>
         statement(assignmentStmt), statement(assignmentStmt),
         Parser<AccEndAtomic>{} / endAccLine))
 
-TYPE_PARSER(construct<OpenACCAtomicConstruct>(Parser<AccAtomicRead>{}) ||
-    construct<OpenACCAtomicConstruct>(Parser<AccAtomicCapture>{}) ||
-    construct<OpenACCAtomicConstruct>(Parser<AccAtomicWrite>{}) ||
-    construct<OpenACCAtomicConstruct>(Parser<AccAtomicUpdate>{}))
+TYPE_PARSER(
+    sourced(construct<OpenACCAtomicConstruct>(Parser<AccAtomicRead>{})) ||
+    sourced(construct<OpenACCAtomicConstruct>(Parser<AccAtomicCapture>{})) ||
+    sourced(construct<OpenACCAtomicConstruct>(Parser<AccAtomicWrite>{})) ||
+    sourced(construct<OpenACCAtomicConstruct>(Parser<AccAtomicUpdate>{})))
 
 // 2.13 Declare constructs
 TYPE_PARSER(sourced(construct<AccDeclarativeDirective>(

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

@@ -294,6 +294,21 @@ void AccStructureChecker::Leave(const parser::OpenACCRoutineConstruct &) {
   dirContext_.pop_back();
 }
 
+void AccStructureChecker::Enter(const parser::OpenACCWaitConstruct &x) {
+  const auto &verbatim{std::get<parser::Verbatim>(x.t)};
+  PushContextAndClauseSets(verbatim.source, llvm::acc::Directive::ACCD_wait);
+}
+void AccStructureChecker::Leave(const parser::OpenACCWaitConstruct &x) {
+  dirContext_.pop_back();
+}
+
+void AccStructureChecker::Enter(const parser::OpenACCAtomicConstruct &x) {
+  PushContextAndClauseSets(x.source, llvm::acc::Directive::ACCD_atomic);
+}
+void AccStructureChecker::Leave(const parser::OpenACCAtomicConstruct &x) {
+  dirContext_.pop_back();
+}
+
 // Clause checkers
 CHECK_REQ_SCALAR_INT_CONSTANT_CLAUSE(Collapse, ACCC_collapse)
 

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

@@ -55,6 +55,10 @@ public:
   void Leave(const parser::OpenACCStandaloneConstruct &);
   void Enter(const parser::OpenACCStandaloneDeclarativeConstruct &);
   void Leave(const parser::OpenACCStandaloneDeclarativeConstruct &);
+  void Enter(const parser::OpenACCWaitConstruct &);
+  void Leave(const parser::OpenACCWaitConstruct &);
+  void Enter(const parser::OpenACCAtomicConstruct &);
+  void Leave(const parser::OpenACCAtomicConstruct &);
 
   // Clauses
   void Leave(const parser::AccClauseList &);

flang/test/Semantics/acc-clause-validity.f90

@@ -6,6 +6,7 @@
 !   2.5.2 Kernels
 !   2.5.3 Serial
 !   2.9 Loop
+!   2.12 Atomic
 !   2.13 Declare
 !   2.14.3 Set
 !   2.14.4 Update
@@ -13,6 +14,7 @@
 !   2.11 Parallel Loop
 !   2.11 Kernels Loop
 !   2.11 Serial Loop
+!   2.16.13 Wait
 
 program openacc_clause_validity
 
@@ -424,6 +426,10 @@ program openacc_clause_validity
   !$acc kernels wait(devnum: 1: queues: 1, 2) async(3)
   !$acc end kernels
 
+  !$acc wait
+
+  !$acc wait async
+
   !$acc wait(1)
   !$acc wait(1, 2)
 
@@ -436,6 +442,46 @@ program openacc_clause_validity
   !$acc wait(devnum: 1: queues: 3)
   !$acc wait(devnum: 1: queues: 3, 4)
 
+  !$acc wait(1) if(.true.)
+
+  !ERROR: At most one IF clause can appear on the WAIT directive
+  !$acc wait(1) if(.true.) if(.false.)
+
+  !$acc wait(1) if(.true.) async
+
+  !$acc wait(1) if(.true.) async(1)
+
+  !ERROR: At most one ASYNC clause can appear on the WAIT directive
+  !$acc wait(1) if(.true.) async(1) async
+
+  !$acc parallel
+  !$acc atomic update
+  c(i) = c(i) + 1
+
+  !$acc atomic update
+  c(i) = c(i) + 1
+  !$acc end atomic
+
+  !$acc atomic write
+  c(i) = 10
+
+  !$acc atomic write
+  c(i) = 10
+  !$acc end atomic
+
+  !$acc atomic read
+  i = c(i)
+
+  !$acc atomic read
+  i = c(i)
+  !$acc end atomic
+
+  !$acc atomic capture
+  c(i) = i
+  i = i + 1
+  !$acc end atomic
+  !$acc end parallel
+
  contains
 
    subroutine sub1(a)