[OpenACC] 'wait' clause for compute construct sema

'wait' takes a few int-exprs (well, a series of async-arguments, but
    those are effectively just an int-expr), plus a pair of tags. This
patch adds the support for this to the AST, and does the appropriate
semantic analysis for them.

Author: erichkeane

clang/include/clang/AST/OpenACCClause.h

@@ -192,6 +192,46 @@ class OpenACCClauseWithExprs : public OpenACCClauseWithParams {
   }
 };
 
+// Represents the 'devnum' and expressions lists for the 'wait' clause.
+class OpenACCWaitClause final
+    : public OpenACCClauseWithExprs,
+      public llvm::TrailingObjects<OpenACCWaitClause, Expr *> {
+  SourceLocation QueuesLoc;
+  OpenACCWaitClause(SourceLocation BeginLoc, SourceLocation LParenLoc,
+                    Expr *DevNumExpr, SourceLocation QueuesLoc,
+                    ArrayRef<Expr *> QueueIdExprs, SourceLocation EndLoc)
+      : OpenACCClauseWithExprs(OpenACCClauseKind::Wait, BeginLoc, LParenLoc,
+                               EndLoc),
+        QueuesLoc(QueuesLoc) {
+    // The first element of the trailing storage is always the devnum expr,
+    // whether it is used or not.
+    std::uninitialized_copy(&DevNumExpr, &DevNumExpr + 1,
+                            getTrailingObjects<Expr *>());
+    std::uninitialized_copy(QueueIdExprs.begin(), QueueIdExprs.end(),
+                            getTrailingObjects<Expr *>() + 1);
+    setExprs(
+        MutableArrayRef(getTrailingObjects<Expr *>(), QueueIdExprs.size() + 1));
+  }
+
+public:
+  static OpenACCWaitClause *Create(const ASTContext &C, SourceLocation BeginLoc,
+                                   SourceLocation LParenLoc, Expr *DevNumExpr,
+                                   SourceLocation QueuesLoc,
+                                   ArrayRef<Expr *> QueueIdExprs,
+                                   SourceLocation EndLoc);
+
+  bool hasQueuesTag() const { return !QueuesLoc.isInvalid(); }
+  SourceLocation getQueuesLoc() const { return QueuesLoc; }
+  bool hasDevNumExpr() const { return getExprs()[0]; }
+  Expr *getDevNumExpr() const { return getExprs()[0]; }
+  llvm::ArrayRef<Expr *> getQueueIdExprs() {
+    return OpenACCClauseWithExprs::getExprs().drop_front();
+  }
+  llvm::ArrayRef<Expr *> getQueueIdExprs() const {
+    return OpenACCClauseWithExprs::getExprs().drop_front();
+  }
+};
+
 class OpenACCNumGangsClause final
     : public OpenACCClauseWithExprs,
       public llvm::TrailingObjects<OpenACCNumGangsClause, Expr *> {

clang/include/clang/Basic/OpenACCClauses.def

@@ -46,6 +46,7 @@ VISIT_CLAUSE(Present)
 VISIT_CLAUSE(Private)
 VISIT_CLAUSE(Self)
 VISIT_CLAUSE(VectorLength)
+VISIT_CLAUSE(Wait)
 
 #undef VISIT_CLAUSE
 #undef CLAUSE_ALIAS

clang/include/clang/Parse/Parser.h

@@ -3632,6 +3632,13 @@ class Parser : public CodeCompletionHandler {
     // Wait constructs, we likely want to put that information in here as well.
   };
 
+  struct OpenACCWaitParseInfo {
+    bool Failed = false;
+    Expr *DevNumExpr = nullptr;
+    SourceLocation QueuesLoc;
+    SmallVector<Expr *> QueueIdExprs;
+  };
+
   /// Represents the 'error' state of parsing an OpenACC Clause, and stores
   /// whether we can continue parsing, or should give up on the directive.
   enum class OpenACCParseCanContinue { Cannot = 0, Can = 1 };
@@ -3674,7 +3681,8 @@ class Parser : public CodeCompletionHandler {
   /// Parses the clause-list for an OpenACC directive.
   SmallVector<OpenACCClause *>
   ParseOpenACCClauseList(OpenACCDirectiveKind DirKind);
-  bool ParseOpenACCWaitArgument(SourceLocation Loc, bool IsDirective);
+  OpenACCWaitParseInfo ParseOpenACCWaitArgument(SourceLocation Loc,
+                                                bool IsDirective);
   /// Parses the clause of the 'bind' argument, which can be a string literal or
   /// an ID expression.
   ExprResult ParseOpenACCBindClauseArgument();

clang/include/clang/Sema/SemaOpenACC.h

@@ -54,8 +54,14 @@ class SemaOpenACC : public SemaBase {
       bool IsZero;
     };
 
+    struct WaitDetails {
+      Expr *DevNumExpr;
+      SourceLocation QueuesLoc;
+      SmallVector<Expr *> QueueIdExprs;
+    };
+
     std::variant<std::monostate, DefaultDetails, ConditionDetails,
-                 IntExprDetails, VarListDetails>
+                 IntExprDetails, VarListDetails, WaitDetails>
         Details = std::monostate{};
 
   public:
@@ -104,14 +110,45 @@ class SemaOpenACC : public SemaBase {
               ClauseKind == OpenACCClauseKind::Async ||
               ClauseKind == OpenACCClauseKind::VectorLength) &&
              "Parsed clause kind does not have a int exprs");
-      //
-      // 'async' has an optional IntExpr, so be tolerant of that.
-      if (ClauseKind == OpenACCClauseKind::Async &&
+
+      // 'async' and 'wait' have an optional IntExpr, so be tolerant of that.
+      if ((ClauseKind == OpenACCClauseKind::Async ||
+           ClauseKind == OpenACCClauseKind::Wait) &&
           std::holds_alternative<std::monostate>(Details))
         return 0;
       return std::get<IntExprDetails>(Details).IntExprs.size();
     }
 
+    SourceLocation getQueuesLoc() const {
+      assert(ClauseKind == OpenACCClauseKind::Wait &&
+             "Parsed clause kind does not have a queues location");
+
+      if (std::holds_alternative<std::monostate>(Details))
+        return SourceLocation{};
+
+      return std::get<WaitDetails>(Details).QueuesLoc;
+    }
+
+    Expr *getDevNumExpr() const {
+      assert(ClauseKind == OpenACCClauseKind::Wait &&
+             "Parsed clause kind does not have a device number expr");
+
+      if (std::holds_alternative<std::monostate>(Details))
+        return nullptr;
+
+      return std::get<WaitDetails>(Details).DevNumExpr;
+    }
+
+    ArrayRef<Expr *> getQueueIdExprs() const {
+      assert(ClauseKind == OpenACCClauseKind::Wait &&
+             "Parsed clause kind does not have a queue id expr list");
+
+      if (std::holds_alternative<std::monostate>(Details))
+        return ArrayRef<Expr *>{std::nullopt};
+
+      return std::get<WaitDetails>(Details).QueueIdExprs;
+    }
+
     ArrayRef<Expr *> getIntExprs() {
       assert((ClauseKind == OpenACCClauseKind::NumGangs ||
               ClauseKind == OpenACCClauseKind::NumWorkers ||
@@ -282,6 +319,13 @@ class SemaOpenACC : public SemaBase {
              "zero: tag only valid on copyout/create");
       Details = VarListDetails{std::move(VarList), IsReadOnly, IsZero};
     }
+
+    void setWaitDetails(Expr *DevNum, SourceLocation QueuesLoc,
+                        llvm::SmallVector<Expr *> &&IntExprs) {
+      assert(ClauseKind == OpenACCClauseKind::Wait &&
+             "Parsed clause kind does not have a wait-details");
+      Details = WaitDetails{DevNum, QueuesLoc, std::move(IntExprs)};
+    }
   };
 
   SemaOpenACC(Sema &S);

clang/include/clang/Serialization/ASTRecordReader.h

@@ -272,6 +272,9 @@ class ASTRecordReader
   /// Read a list of Exprs used for a var-list.
   llvm::SmallVector<Expr *> readOpenACCVarList();
 
+  /// Read a list of Exprs used for a int-expr-list.
+  llvm::SmallVector<Expr *> readOpenACCIntExprList();
+
   /// Read an OpenACC clause, advancing Idx.
   OpenACCClause *readOpenACCClause();
 

clang/include/clang/Serialization/ASTRecordWriter.h

@@ -296,6 +296,8 @@ class ASTRecordWriter
 
   void writeOpenACCVarList(const OpenACCClauseWithVarList *C);
 
+  void writeOpenACCIntExprList(ArrayRef<Expr *> Exprs);
+
   /// Writes out a single OpenACC Clause.
   void writeOpenACCClause(const OpenACCClause *C);
 

clang/lib/AST/OpenACCClause.cpp

@@ -147,6 +147,18 @@ OpenACCAsyncClause *OpenACCAsyncClause::Create(const ASTContext &C,
   return new (Mem) OpenACCAsyncClause(BeginLoc, LParenLoc, IntExpr, EndLoc);
 }
 
+OpenACCWaitClause *OpenACCWaitClause::Create(
+    const ASTContext &C, SourceLocation BeginLoc, SourceLocation LParenLoc,
+    Expr *DevNumExpr, SourceLocation QueuesLoc, ArrayRef<Expr *> QueueIdExprs,
+    SourceLocation EndLoc) {
+  // Allocates enough room in trailing storage for all the int-exprs, plus a
+  // placeholder for the devnum.
+  void *Mem = C.Allocate(
+      OpenACCWaitClause::totalSizeToAlloc<Expr *>(QueueIdExprs.size() + 1));
+  return new (Mem) OpenACCWaitClause(BeginLoc, LParenLoc, DevNumExpr, QueuesLoc,
+                                     QueueIdExprs, EndLoc);
+}
+
 OpenACCNumGangsClause *OpenACCNumGangsClause::Create(const ASTContext &C,
                                                      SourceLocation BeginLoc,
                                                      SourceLocation LParenLoc,
@@ -393,3 +405,22 @@ void OpenACCClausePrinter::VisitCreateClause(const OpenACCCreateClause &C) {
                         [&](const Expr *E) { printExpr(E); });
   OS << ")";
 }
+
+void OpenACCClausePrinter::VisitWaitClause(const OpenACCWaitClause &C) {
+  OS << "wait";
+  if (!C.getLParenLoc().isInvalid()) {
+    OS << "(";
+    if (C.hasDevNumExpr()) {
+      OS << "devnum: ";
+      printExpr(C.getDevNumExpr());
+      OS << " : ";
+    }
+
+    if (C.hasQueuesTag())
+      OS << "queues: ";
+
+    llvm::interleaveComma(C.getQueueIdExprs(), OS,
+                          [&](const Expr *E) { printExpr(E); });
+    OS << ")";
+  }
+}

clang/lib/AST/StmtProfile.cpp

@@ -2578,6 +2578,13 @@ void OpenACCClauseProfiler::VisitAsyncClause(const OpenACCAsyncClause &Clause) {
   if (Clause.hasIntExpr())
     Profiler.VisitStmt(Clause.getIntExpr());
 }
+
+void OpenACCClauseProfiler::VisitWaitClause(const OpenACCWaitClause &Clause) {
+  if (Clause.hasDevNumExpr())
+    Profiler.VisitStmt(Clause.getDevNumExpr());
+  for (auto *E : Clause.getQueueIdExprs())
+    Profiler.VisitStmt(E);
+}
 } // namespace
 
 void StmtProfiler::VisitOpenACCComputeConstruct(

clang/lib/AST/TextNodeDumper.cpp

@@ -437,6 +437,13 @@ void TextNodeDumper::Visit(const OpenACCClause *C) {
       if (cast<OpenACCCreateClause>(C)->isZero())
         OS << " : zero";
       break;
+    case OpenACCClauseKind::Wait:
+      OS << " clause";
+      if (cast<OpenACCWaitClause>(C)->hasDevNumExpr())
+        OS << " has devnum";
+      if (cast<OpenACCWaitClause>(C)->hasQueuesTag())
+        OS << " has queues tag";
+      break;
     default:
       // Nothing to do here.
       break;

clang/lib/Parse/ParseOpenACC.cpp

@@ -867,7 +867,6 @@ Parser::OpenACCClauseParseResult Parser::ParseOpenACCClauseParams(
   SemaOpenACC::OpenACCParsedClause ParsedClause(DirKind, ClauseKind, ClauseLoc);
 
   if (ClauseHasRequiredParens(DirKind, ClauseKind)) {
-    ParsedClause.setLParenLoc(getCurToken().getLocation());
     if (Parens.expectAndConsume()) {
       // We are missing a paren, so assume that the person just forgot the
       // parameter.  Return 'false' so we try to continue on and parse the next
@@ -876,6 +875,7 @@ Parser::OpenACCClauseParseResult Parser::ParseOpenACCClauseParams(
                 Parser::StopBeforeMatch);
       return OpenACCCanContinue();
     }
+    ParsedClause.setLParenLoc(Parens.getOpenLocation());
 
     switch (ClauseKind) {
     case OpenACCClauseKind::Default: {
@@ -1048,8 +1048,8 @@ Parser::OpenACCClauseParseResult Parser::ParseOpenACCClauseParams(
       return OpenACCCannotContinue();
 
   } else if (ClauseHasOptionalParens(DirKind, ClauseKind)) {
-    ParsedClause.setLParenLoc(getCurToken().getLocation());
     if (!Parens.consumeOpen()) {
+      ParsedClause.setLParenLoc(Parens.getOpenLocation());
       switch (ClauseKind) {
       case OpenACCClauseKind::Self: {
         assert(DirKind != OpenACCDirectiveKind::Update);
@@ -1099,13 +1099,19 @@ Parser::OpenACCClauseParseResult Parser::ParseOpenACCClauseParams(
           return OpenACCCanContinue();
         }
         break;
-      case OpenACCClauseKind::Wait:
-        if (ParseOpenACCWaitArgument(ClauseLoc,
-                                     /*IsDirective=*/false)) {
+      case OpenACCClauseKind::Wait: {
+        OpenACCWaitParseInfo Info =
+            ParseOpenACCWaitArgument(ClauseLoc,
+                                     /*IsDirective=*/false);
+        if (Info.Failed) {
           Parens.skipToEnd();
           return OpenACCCanContinue();
         }
+
+        ParsedClause.setWaitDetails(Info.DevNumExpr, Info.QueuesLoc,
+                                    std::move(Info.QueueIdExprs));
         break;
+      }
       default:
         llvm_unreachable("Not an optional parens type?");
       }
@@ -1139,7 +1145,9 @@ Parser::ParseOpenACCAsyncArgument(OpenACCDirectiveKind DK, OpenACCClauseKind CK,
 /// In this section and throughout the specification, the term wait-argument
 /// means:
 /// [ devnum : int-expr : ] [ queues : ] async-argument-list
-bool Parser::ParseOpenACCWaitArgument(SourceLocation Loc, bool IsDirective) {
+Parser::OpenACCWaitParseInfo
+Parser::ParseOpenACCWaitArgument(SourceLocation Loc, bool IsDirective) {
+  OpenACCWaitParseInfo Result;
   // [devnum : int-expr : ]
   if (isOpenACCSpecialToken(OpenACCSpecialTokenKind::DevNum, Tok) &&
       NextToken().is(tok::colon)) {
@@ -1153,18 +1161,25 @@ bool Parser::ParseOpenACCWaitArgument(SourceLocation Loc, bool IsDirective) {
                     : OpenACCDirectiveKind::Invalid,
         IsDirective ? OpenACCClauseKind::Invalid : OpenACCClauseKind::Wait,
         Loc);
-    if (Res.first.isInvalid() && Res.second == OpenACCParseCanContinue::Cannot)
-      return true;
+    if (Res.first.isInvalid() &&
+        Res.second == OpenACCParseCanContinue::Cannot) {
+      Result.Failed = true;
+      return Result;
+    }
 
-    if (ExpectAndConsume(tok::colon))
-      return true;
+    if (ExpectAndConsume(tok::colon)) {
+      Result.Failed = true;
+      return Result;
+    }
+
+    Result.DevNumExpr = Res.first.get();
   }
 
   // [ queues : ]
   if (isOpenACCSpecialToken(OpenACCSpecialTokenKind::Queues, Tok) &&
       NextToken().is(tok::colon)) {
     // Consume queues.
-    ConsumeToken();
+    Result.QueuesLoc = ConsumeToken();
     // Consume colon.
     ConsumeToken();
   }
@@ -1176,8 +1191,10 @@ bool Parser::ParseOpenACCWaitArgument(SourceLocation Loc, bool IsDirective) {
   bool FirstArg = true;
   while (!getCurToken().isOneOf(tok::r_paren, tok::annot_pragma_openacc_end)) {
     if (!FirstArg) {
-      if (ExpectAndConsume(tok::comma))
-        return true;
+      if (ExpectAndConsume(tok::comma)) {
+        Result.Failed = true;
+        return Result;
+      }
     }
     FirstArg = false;
 
@@ -1187,11 +1204,16 @@ bool Parser::ParseOpenACCWaitArgument(SourceLocation Loc, bool IsDirective) {
         IsDirective ? OpenACCClauseKind::Invalid : OpenACCClauseKind::Wait,
         Loc);
 
-    if (Res.first.isInvalid() && Res.second == OpenACCParseCanContinue::Cannot)
-      return true;
+    if (Res.first.isInvalid() &&
+        Res.second == OpenACCParseCanContinue::Cannot) {
+      Result.Failed = true;
+      return Result;
+    }
+
+    Result.QueueIdExprs.push_back(Res.first.get());
   }
 
-  return false;
+  return Result;
 }
 
 ExprResult Parser::ParseOpenACCIDExpression() {
@@ -1360,7 +1382,7 @@ Parser::OpenACCDirectiveParseInfo Parser::ParseOpenACCDirective() {
       break;
     case OpenACCDirectiveKind::Wait:
       // OpenACC has an optional paren-wrapped 'wait-argument'.
-      if (ParseOpenACCWaitArgument(StartLoc, /*IsDirective=*/true))
+      if (ParseOpenACCWaitArgument(StartLoc, /*IsDirective=*/true).Failed)
         T.skipToEnd();
       else
         T.consumeClose();