[OpenMP] Support runtime user conditions in metadirective

[ggeorgakoudis]

Summary

This revision adds support for dynamic user conditions in OpenMP metadirectives resolvable at runtime. Parsing identifies when clauses of variants (including the default clause) and matches against the OpenMP context to forward applicable when clauses for code generation. Code generation generates dynamic variants checking runtime user conditions and the best match of the statically resolvable condition variants to fall back.

[ggeorgakoudis]

Fix tentative parsing reverting
Fix device codegen for metadirectives

[abidmalikwaterloo]

Does anyone implement/take care of the "otherwise" clause? If I am correct, OpenMP 5.2 supports otherwise instead of the default for meta directive.

[abidmalikwaterloo]

clang/lib/CodeGen/CGStmtOpenMP.cpp

1852

This approach is valid for static conditions. How you are going to take care of multiple dynamic/runtime conditions?

[abidmalikwaterloo]

clang/lib/CodeGen/CGStmtOpenMP.cpp

1793

Why are you skipping the AST building part? Any special significance of this approach?

[ggeorgakoudis]

clang/lib/CodeGen/CGStmtOpenMP.cpp

1793

I don't understand the question, what do you mean by that? Emit generates code to handle dynamic and static conditions based on the when clauses

1852

See GenerateCond that creates code for dynamic conditions. If the when clause has dynamically evaluated conditions it generates basicblocks and control logic.

[abidmalikwaterloo]

clang/lib/CodeGen/CGStmtOpenMP.cpp

1852

Why get a score for static when clauses for the best clause? Is it according to OpenMP specification?

[abidmalikwaterloo]

Are you going to work on it? Just confirming. I am working on D122255. The code generation part for the dynamic condition is useful. However, my understanding is that it is not according to the OpenMP 5.2 specification. Correct me if I am wrong.

[ggeorgakoudis]

In D120573#3415517, @abidmalikwaterloo wrote:

Are you going to work on it? Just confirming. I am working on D122255. The code generation part for the dynamic condition is useful. However, my understanding is that it is not according to the OpenMP 5.2 specification. Correct me if I am wrong.

I'm happy if you'd like to take over in your patch. I'm not sure what is the issue with the OpenMP 5.2 but feel free to re-use/modify parts of this patch. One thing that is problematic with this patch is the need to re-set the tokens in the lexer/parser to their unannotated versions when tentative parsing. The code in the patch is not fault-proof but I have an update for that that I can post, if that's useful. In D122255 I see you're taking the route of parsing the associated statement which has its own pitfalls but maybe that's the best approach.

clang/lib/CodeGen/CGStmtOpenMP.cpp

1852

Because if no runtime conditions evaluate to true then it falls back to the static when with the highest score. I'm not sure about the OpenMP specification. Do you have an example? Could you point me to the related part in the specification it fails to conform to?

clang/include/clang/AST/OpenMPClause.h

const_child_range used_children() const {
     return const_child_range(const_child_iterator(), const_child_iterator());
   }
 
   static bool classof(const OMPClause *T) {
     return T->getClauseKind() == llvm::omp::OMPC_bind;
   }
 };
 
+/// This represents 'when' clause in the '#pragma omp metadirective'
+/// directive.
+///
+/// \code
+/// #pragma omp metadirective when(user={condition(N<10)}: parallel)
+/// \endcode
+/// In this example directive '#pragma omp metadirective' has simple 'when'
+/// clause with user defined condition.
+class OMPTraitInfo;
+class OMPWhenClause final : public OMPClause {
+  friend class OMPClauseReader;
+
+  OMPTraitInfo *TI;
+  Stmt *Directive;
+
+  /// Location of '('.
+  SourceLocation LParenLoc;
+
+public:
+  /// Build 'when' clause with arguments \a T for traits, \a D for the
+  /// associated directive.
+  ///
+  /// \param T TraitInfo containing information about the context selector
+  /// \param D The statement associated with the when clause
+  /// \param StartLoc Starting location of the clause.
+  /// \param LParenLoc Location of '('.
+  /// \param EndLoc Ending location of the clause.
+  OMPWhenClause(OMPTraitInfo &T, Stmt *D, SourceLocation StartLoc,
+                SourceLocation LParenLoc, SourceLocation EndLoc)
+      : OMPClause(llvm::omp::OMPC_when, StartLoc, EndLoc), TI(&T), Directive(D),
+        LParenLoc(LParenLoc) {}
+
+  /// Build an empty clause.
+  OMPWhenClause()
+      : OMPClause(llvm::omp::OMPC_when, SourceLocation(), SourceLocation()) {}
+
+  /// Sets the location of '('.
+  void setLParenLoc(SourceLocation Loc) { LParenLoc = Loc; }
+
+  /// Returns the location of '('.
+  SourceLocation getLParenLoc() const { return LParenLoc; }
+
+  /// Returns the associated OpenMP directive.
+  Stmt *getDirective() const { return Directive; }
+
+  /// Set the associated OpenMP directive.
+  void setDirective(Stmt *S) { Directive = S; }
+
+  /// Returns the OMPTraitInfo
+  OMPTraitInfo &getTraitInfo() const { return *TI; }
+
+  /// Set the OMPTraitInfo
+  void setTraitInfo(OMPTraitInfo *T) { TI = T; }
+
+  child_range children() {
+    return child_range(child_iterator(), child_iterator());
+  }
+
+  const_child_range children() const {
+    return const_child_range(const_child_iterator(), const_child_iterator());
+  }
+
+  child_range used_children() {
+    return child_range(child_iterator(), child_iterator());
+  }
+  const_child_range used_children() const {
+    return const_child_range(const_child_iterator(), const_child_iterator());
+  }
+
+  static bool classof(const OMPClause *T) {
+    return T->getClauseKind() == llvm::omp::OMPC_when;
+  }
+};
+
 /// This class implements a simple visitor for OMPClause
 /// subclasses.
 template<class ImplClass, template <typename> class Ptr, typename RetTy>
 class OMPClauseVisitorBase {
 public:
 #define PTR(CLASS) Ptr<CLASS>
 #define DISPATCH(CLASS) \
   return static_cast<ImplClass*>(this)->Visit##CLASS(static_cast<PTR(CLASS)>(S))

clang/include/clang/AST/RecursiveASTVisitor.h

bool RecursiveASTVisitor<Derived>::VisitOMPFilterClause(OMPFilterClause *C) {
   return true;
 }
 
 template <typename Derived>
 bool RecursiveASTVisitor<Derived>::VisitOMPBindClause(OMPBindClause *C) {
   return true;
 }
 
+template <typename Derived>
+bool RecursiveASTVisitor<Derived>::VisitOMPWhenClause(OMPWhenClause *C) {
+  for (const OMPTraitSet &Set : C->getTraitInfo().Sets) {
+    for (const OMPTraitSelector &Selector : Set.Selectors) {
+      if (Selector.Kind == llvm::omp::TraitSelector::user_condition &&
+          Selector.ScoreOrCondition)
+        TRY_TO(TraverseStmt(Selector.ScoreOrCondition));
+    }
+  }
+  return true;
+}
+
 // FIXME: look at the following tricky-seeming exprs to see if we
 // need to recurse on anything.  These are ones that have methods
 // returning decls or qualtypes or nestednamespecifier -- though I'm
 // not sure if they own them -- or just seemed very complicated, or
 // had lots of sub-types to explore.
 //
 // VisitOverloadExpr and its children: recurse on template args? etc?
 

clang/include/clang/AST/StmtOpenMP.h

 /// #pragma omp metadirective when(user={condition(N>10)}: parallel for)
 /// \endcode
 /// In this example directive '#pragma omp metadirective' has clauses 'when'
 /// with a dynamic user condition to check if a variable 'N > 10'
 ///
 class OMPMetaDirective final : public OMPExecutableDirective {
   friend class ASTStmtReader;
   friend class OMPExecutableDirective;
-  Stmt *IfStmt;
 
   OMPMetaDirective(SourceLocation StartLoc, SourceLocation EndLoc)
       : OMPExecutableDirective(OMPMetaDirectiveClass,
                                llvm::omp::OMPD_metadirective, StartLoc,
                                EndLoc) {}
   explicit OMPMetaDirective()
       : OMPExecutableDirective(OMPMetaDirectiveClass,
                                llvm::omp::OMPD_metadirective, SourceLocation(),
                                SourceLocation()) {}
 
-  void setIfStmt(Stmt *S) { IfStmt = S; }
-
 public:
   static OMPMetaDirective *Create(const ASTContext &C, SourceLocation StartLoc,
                                   SourceLocation EndLoc,
-                                  ArrayRef<OMPClause *> Clauses,
-                                  Stmt *AssociatedStmt, Stmt *IfStmt);
+                                  ArrayRef<OMPClause *> Clauses);
+
   static OMPMetaDirective *CreateEmpty(const ASTContext &C, unsigned NumClauses,
                                        EmptyShell);
-  Stmt *getIfStmt() const { return IfStmt; }
 
   static bool classof(const Stmt *T) {
     return T->getStmtClass() == OMPMetaDirectiveClass;
   }
 };
 
 /// This represents '#pragma omp loop' directive.
 ///

clang/include/clang/Lex/Preprocessor.h

private:
   PreprocessingRecord *Record = nullptr;
 
   /// Cached tokens state.
   using CachedTokensTy = SmallVector<Token, 1>;
 
   /// Cached tokens are stored here when we do backtracking or
   /// lookahead. They are "lexed" by the CachingLex() method.
   CachedTokensTy CachedTokens;
+  CachedTokensTy UnannotatedCachedTokens;
 
   /// The position of the cached token that CachingLex() should
   /// "lex" next.
   ///
   /// If it points beyond the CachedTokens vector, it means that a normal
   /// Lex() should be invoked.
   CachedTokensTy::size_type CachedLexPos = 0;
 
public:
   ///
   void EnableBacktrackAtThisPos();
 
   /// Disable the last EnableBacktrackAtThisPos call.
   void CommitBacktrackedTokens();
 
   /// Make Preprocessor re-lex the tokens that were lexed since
   /// EnableBacktrackAtThisPos() was previously called.
-  void Backtrack();
+  void Backtrack(bool RevertAnnotations);
 
   /// True if EnableBacktrackAtThisPos() was called and
   /// caching of tokens is on.
   bool isBacktrackEnabled() const { return !BacktrackPositions.empty(); }
 
   /// Lex the next token for this preprocessor.
   void Lex(Token &Result);
 

clang/include/clang/Parse/Parser.h

public:
     }
     void Commit() {
       assert(isActive && "Parsing action was finished!");
       P.TentativelyDeclaredIdentifiers.resize(
           PrevTentativelyDeclaredIdentifierCount);
       P.PP.CommitBacktrackedTokens();
       isActive = false;
     }
-    void Revert() {
+    void Revert(bool ReverseAnnotations = false) {
       assert(isActive && "Parsing action was finished!");
-      P.PP.Backtrack();
+      P.PP.Backtrack(ReverseAnnotations);
       P.PreferredType = PrevPreferredType;
       P.Tok = PrevTok;
       P.TentativelyDeclaredIdentifiers.resize(
           PrevTentativelyDeclaredIdentifierCount);
       P.ParenCount = PrevParenCount;
       P.BracketCount = PrevBracketCount;
       P.BraceCount = PrevBraceCount;
       isActive = false;

clang/include/clang/Sema/Sema.h

public:
   /// mark loop control variable, used in \p Init for loop initialization, as
   /// private by default.
   /// \param Init First part of the for loop.
   void ActOnOpenMPLoopInitialization(SourceLocation ForLoc, Stmt *Init);
 
   /// Called on well-formed '\#pragma omp metadirective' after parsing
   /// of the  associated statement.
   StmtResult ActOnOpenMPMetaDirective(ArrayRef<OMPClause *> Clauses,
-                                      Stmt *AStmt, SourceLocation StartLoc,
+                                      SourceLocation StartLoc,
                                       SourceLocation EndLoc);
 
   // OpenMP directives and clauses.
   /// Called on correct id-expression from the '#pragma omp
   /// threadprivate'.
   ExprResult ActOnOpenMPIdExpression(Scope *CurScope, CXXScopeSpec &ScopeSpec,
                                      const DeclarationNameInfo &Id,
                                      OpenMPDirectiveKind Kind);
public:
 
   OMPClause *ActOnOpenMPSimpleClause(OpenMPClauseKind Kind,
                                      unsigned Argument,
                                      SourceLocation ArgumentLoc,
                                      SourceLocation StartLoc,
                                      SourceLocation LParenLoc,
                                      SourceLocation EndLoc);
   /// Called on well-formed 'when' clause.
-  OMPClause *ActOnOpenMPWhenClause(OMPTraitInfo &TI, SourceLocation StartLoc,
+  OMPClause *ActOnOpenMPWhenClause(OMPTraitInfo &TI, StmtResult Directive,
+                                   SourceLocation StartLoc,
                                    SourceLocation LParenLoc,
                                    SourceLocation EndLoc);
   /// Called on well-formed 'default' clause.
   OMPClause *ActOnOpenMPDefaultClause(llvm::omp::DefaultKind Kind,
                                       SourceLocation KindLoc,
                                       SourceLocation StartLoc,
                                       SourceLocation LParenLoc,
                                       SourceLocation EndLoc);

clang/lib/AST/OpenMPClause.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "clang/AST/OpenMPClause.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclOpenMP.h"
+#include "clang/AST/StmtOpenMP.h"
 #include "clang/Basic/LLVM.h"
 #include "clang/Basic/OpenMPKinds.h"
 #include "clang/Basic/TargetInfo.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include <algorithm>
 #include <cassert>

 }
 
 void OMPClausePrinter::VisitOMPBindClause(OMPBindClause *Node) {
   OS << "bind("
      << getOpenMPSimpleClauseTypeName(OMPC_bind, unsigned(Node->getBindKind()))
      << ")";
 }
 
+void OMPClausePrinter::VisitOMPWhenClause(OMPWhenClause *Node) {
+  OMPTraitInfo &TI = Node->getTraitInfo();
+  if (TI.Sets.empty())
+    OS << "default(";
+  else
+    OS << "when(";
+  TI.print(OS, Policy);
+  Stmt *S = Node->getDirective();
+  if (S) {
+    OS << ":";
+    OMPExecutableDirective *D = cast<OMPExecutableDirective>(S);
+    auto DKind = D->getDirectiveKind();
+    OS << getOpenMPDirectiveName(DKind);
+
+    OMPClausePrinter Printer(OS, Policy);
+    ArrayRef<OMPClause *> Clauses = D->clauses();
+    for (auto *Clause : Clauses)
+      if (Clause && !Clause->isImplicit()) {
+        OS << ' ';
+        Printer.Visit(Clause);
+      }
+  }
+  OS << ")";
+  return;
+}
+
 void OMPTraitInfo::getAsVariantMatchInfo(ASTContext &ASTCtx,
                                          VariantMatchInfo &VMI) const {
   for (const OMPTraitSet &Set : Sets) {
     for (const OMPTraitSelector &Selector : Set.Selectors) {
 
       // User conditions are special as we evaluate the condition here.
       if (Selector.Kind == TraitSelector::user_condition) {
         assert(Selector.ScoreOrCondition &&
                "Ill-formed user condition, expected condition expression!");
         assert(Selector.Properties.size() == 1 &&
                Selector.Properties.front().Kind ==
                    TraitProperty::user_condition_unknown &&
                "Ill-formed user condition, expected unknown trait property!");
 
+        // If Condition is statically resolvable add it as a trait, otherwise
+        // do nothing since codegen will generate dynamic conditions.
         if (Optional<APSInt> CondVal =
                 Selector.ScoreOrCondition->getIntegerConstantExpr(ASTCtx))
           VMI.addTrait(CondVal->isZero() ? TraitProperty::user_condition_false
                                          : TraitProperty::user_condition_true,
                        "<condition>");
-        else
-          VMI.addTrait(TraitProperty::user_condition_false, "<condition>");
         continue;
       }
 
       Optional<llvm::APSInt> Score;
       llvm::APInt *ScorePtr = nullptr;
       if (Selector.ScoreOrCondition) {
         if ((Score = Selector.ScoreOrCondition->getIntegerConstantExpr(ASTCtx)))
           ScorePtr = &*Score;

clang/lib/AST/StmtOpenMP.cpp

void OMPLoopDirective::setFinalsConditions(ArrayRef<Expr *> A) {
   assert(A.size() == getLoopsNumber() &&
          "Number of finals conditions is not the same as the collapsed number");
   llvm::copy(A, getFinalsConditions().begin());
 }
 
 OMPMetaDirective *OMPMetaDirective::Create(const ASTContext &C,
                                            SourceLocation StartLoc,
                                            SourceLocation EndLoc,
-                                           ArrayRef<OMPClause *> Clauses,
-                                           Stmt *AssociatedStmt, Stmt *IfStmt) {
-  auto *Dir = createDirective<OMPMetaDirective>(
-      C, Clauses, AssociatedStmt, /*NumChildren=*/1, StartLoc, EndLoc);
-  Dir->setIfStmt(IfStmt);
+                                           ArrayRef<OMPClause *> Clauses) {
+  auto *Dir =
+      createDirective<OMPMetaDirective>(C, Clauses, /*AssociatedStmt*/ nullptr,
+                                        /*NumChildren=*/0, StartLoc, EndLoc);
   return Dir;
 }
 
 OMPMetaDirective *OMPMetaDirective::CreateEmpty(const ASTContext &C,
                                                 unsigned NumClauses,
                                                 EmptyShell) {
   return createEmptyDirective<OMPMetaDirective>(C, NumClauses,
                                                 /*HasAssociatedStmt=*/true,

clang/lib/AST/StmtProfile.cpp

 void OMPClauseProfiler::VisitOMPAffinityClause(const OMPAffinityClause *C) {
   if (const Expr *Modifier = C->getModifier())
     Profiler->VisitStmt(Modifier);
   for (const Expr *E : C->varlists())
     Profiler->VisitStmt(E);
 }
 void OMPClauseProfiler::VisitOMPOrderClause(const OMPOrderClause *C) {}
 void OMPClauseProfiler::VisitOMPBindClause(const OMPBindClause *C) {}
+
+void OMPClauseProfiler::VisitOMPWhenClause(const OMPWhenClause *C) {
+  if (C->getDirective())
+    Profiler->VisitStmt(C->getDirective());
+}
 } // namespace
 
 void
 StmtProfiler::VisitOMPExecutableDirective(const OMPExecutableDirective *S) {
   VisitStmt(S);
   OMPClauseProfiler P(this);
   ArrayRef<OMPClause *> Clauses = S->clauses();
   for (ArrayRef<OMPClause *>::iterator I = Clauses.begin(), E = Clauses.end();

clang/lib/CodeGen/CGOpenMPRuntime.cpp

if (RequiresDeviceCodegen) {
     case OMPD_unknown:
     default:
       llvm_unreachable("Unknown target directive for OpenMP device codegen.");
     }
     return;
   }
 
   if (const auto *E = dyn_cast<OMPExecutableDirective>(S)) {
+    if (E->getDirectiveKind() == OMPD_metadirective) {
+      for (const auto *C : E->getClausesOfKind<OMPWhenClause>())
+        if (C->getDirective())
+          scanForTargetRegionsFunctions(C->getDirective(), ParentName);
+    }
+
     if (!E->hasAssociatedStmt() || !E->getAssociatedStmt())
       return;
 
     scanForTargetRegionsFunctions(E->getRawStmt(), ParentName);
     return;
   }
 
   // If this is a lambda function, look into its body.

clang/lib/CodeGen/CGStmtOpenMP.cpp

 #include "clang/AST/Attr.h"
 #include "clang/AST/DeclOpenMP.h"
 #include "clang/AST/OpenMPClause.h"
 #include "clang/AST/Stmt.h"
 #include "clang/AST/StmtOpenMP.h"
 #include "clang/AST/StmtVisitor.h"
 #include "clang/Basic/OpenMPKinds.h"
 #include "clang/Basic/PrettyStackTrace.h"
+#include "clang/Parse/ParseDiagnostic.h"
 #include "llvm/BinaryFormat/Dwarf.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/Support/AtomicOrdering.h"
emitCommonOMPParallelDirective(*this, S, OMPD_parallel, CodeGen,
                                    emitEmptyBoundParameters);
     emitPostUpdateForReductionClause(*this, S,
                                      [](CodeGenFunction &) { return nullptr; });
   }
   // Check for outer lastprivate conditional update.
   checkForLastprivateConditionalUpdate(*this, S);
 }
 
-void CodeGenFunction::EmitOMPMetaDirective(const OMPMetaDirective &S) {
-  EmitStmt(S.getIfStmt());
+void CodeGenFunction::EmitOMPMetaDirective(const OMPMetaDirective &D) {
+  llvm::BasicBlock *AfterBlock =
+      createBasicBlock("omp.meta.user.condition.after");
+
+  SmallVector<const OMPWhenClause *, 4> StaticWhenClauses;
+  SmallVector<VariantMatchInfo, 4> StaticVMIs;
+
+  for (auto *C : D.getClausesOfKind<OMPWhenClause>()) {
+    OMPTraitInfo &TI = C->getTraitInfo();
+
+    llvm::BasicBlock *ExitBlock =
+        createBasicBlock("omp.meta.user.condition.exit");
+
+    // Emit code to generate a dynamic condition, returns true if there is a
+    // condition, false otherwise.
+    auto GenerateCond = [&](Expr *&E, bool IsScore) {
+      if (IsScore)
+        return false;
+
+      // Do not emit code if the expression is statically resolvable, will be
+      // handled as a static when clause.
+      if (E->getIntegerConstantExpr(getContext()))
+        return false;
+
+      llvm::BasicBlock *TrueBlock = createBasicBlock("omp.meta.user.condition");
+      EmitBranchOnBoolExpr(E, TrueBlock, ExitBlock,
+                           getProfileCount(C->getDirective()));
+      EmitBlock(TrueBlock);
+      EmitStmt(C->getDirective());
+      EmitBranch(AfterBlock);
+
+      return true;
+    };
+
+    // If there is no dynamic condition for the clause then add it to the
+    // static clauses and resolve later the best match to generate code for.
+    // This also handles the default clause.
+    if (!TI.anyScoreOrCondition(GenerateCond)) {
+      StaticWhenClauses.push_back(C);
+      VariantMatchInfo VMI;
+      TI.getAsVariantMatchInfo(getContext(), VMI);
+      StaticVMIs.push_back(VMI);
+    } else
+      EmitBlock(ExitBlock);
+  }
+
+  // Emit code for static clauses, if any.
+  if (!StaticWhenClauses.empty()) {
+    std::function<void(StringRef)> DiagUnknownTrait = [&](StringRef ISATrait) {
+      CGM.getDiags().Report(D.getBeginLoc(),
+                            diag::warn_unknown_declare_variant_isa_trait)
+          << ISATrait;
+    };
+
+    TargetOMPContext OMPCtx(
+        getContext(), std::move(DiagUnknownTrait),
+        /* CurrentFunctionDecl */ nullptr,
+        /* ConstructTraits */ ArrayRef<llvm::omp::TraitProperty>());
+
+    int BestIdx = getBestVariantMatchForContext(StaticVMIs, OMPCtx);
+
+    EmitStmt(StaticWhenClauses[BestIdx]->getDirective());
+    EmitBranch(AfterBlock);
+  }
+
+  EmitBlock(AfterBlock);
 }
 
 namespace {
 /// RAII to handle scopes for loop transformation directives.
 class OMPTransformDirectiveScopeRAII {
   OMPLoopScope *Scope = nullptr;
   CodeGenFunction::CGCapturedStmtInfo *CGSI = nullptr;
   CodeGenFunction::CGCapturedStmtRAII *CapInfoRAII = nullptr;

clang/lib/Lex/PPCaching.cpp

 void Preprocessor::CommitBacktrackedTokens() {
   assert(!BacktrackPositions.empty()
          && "EnableBacktrackAtThisPos was not called!");
   BacktrackPositions.pop_back();
 }
 
 // Make Preprocessor re-lex the tokens that were lexed since
 // EnableBacktrackAtThisPos() was previously called.
-void Preprocessor::Backtrack() {
-  assert(!BacktrackPositions.empty()
-         && "EnableBacktrackAtThisPos was not called!");
+void Preprocessor::Backtrack(bool ReverseAnnotations = false) {
+  assert(!BacktrackPositions.empty() &&
+         "EnableBacktrackAtThisPos was not called!");
   CachedLexPos = BacktrackPositions.back();
   BacktrackPositions.pop_back();
   recomputeCurLexerKind();
+
+  if (ReverseAnnotations) {
+    CachedTokens.erase(CachedTokens.begin() + CachedLexPos, CachedTokens.end());
+    CachedTokens.append(UnannotatedCachedTokens.begin() + CachedLexPos,
+                        UnannotatedCachedTokens.end());
+  }
 }
 
 void Preprocessor::CachingLex(Token &Result) {
   if (!InCachingLexMode())
     return;
 
   // The assert in EnterCachingLexMode should prevent this from happening.
   assert(LexLevel == 1 &&
          "should not use token caching within the preprocessor");
 
   if (CachedLexPos < CachedTokens.size()) {
     Result = CachedTokens[CachedLexPos++];
     Result.setFlag(Token::IsReinjected);
     return;
   }
 
   ExitCachingLexMode();
   Lex(Result);
 
   if (isBacktrackEnabled()) {
     // Cache the lexed token.
     EnterCachingLexModeUnchecked();
     CachedTokens.push_back(Result);
+    UnannotatedCachedTokens.push_back(Result);
     ++CachedLexPos;
     return;
   }
 
   if (CachedLexPos < CachedTokens.size()) {
     EnterCachingLexModeUnchecked();
   } else {
     // All cached tokens were consumed.
     CachedTokens.clear();
+    UnannotatedCachedTokens.clear();
     CachedLexPos = 0;
   }
 }
 
 void Preprocessor::EnterCachingLexMode() {
   // The caching layer sits on top of all the other lexers, so it's incorrect
   // to cache tokens while inside a nested lex action. The cached tokens would
   // be retained after returning to the enclosing lex action and, at best,
void Preprocessor::EnterCachingLexModeUnchecked() {
   CurLexerKind = CLK_CachingLexer;
 }
 
 
 const Token &Preprocessor::PeekAhead(unsigned N) {
   assert(CachedLexPos + N > CachedTokens.size() && "Confused caching.");
   ExitCachingLexMode();
   for (size_t C = CachedLexPos + N - CachedTokens.size(); C > 0; --C) {
-    CachedTokens.push_back(Token());
-    Lex(CachedTokens.back());
+    Token Result;
+    Lex(Result);
+    CachedTokens.push_back(Result);
+    UnannotatedCachedTokens.push_back(Result);
   }
   EnterCachingLexMode();
   return CachedTokens.back();
 }
 
 void Preprocessor::AnnotatePreviousCachedTokens(const Token &Tok) {
   assert(Tok.isAnnotation() && "Expected annotation token");
   assert(CachedLexPos != 0 && "Expected to have some cached tokens");

clang/lib/Parse/ParseOpenMP.cpp

Parser::ParseOpenMPDeclarativeOrExecutableDirective(ParsedStmtContext StmtCtx) {
   switch (DKind) {
   case OMPD_metadirective: {
     ConsumeToken();
     SmallVector<VariantMatchInfo, 4> VMIs;
 
     // First iteration of parsing all clauses of metadirective.
     // This iteration only parses and collects all context selector ignoring the
     // associated directives.
-    TentativeParsingAction TPA(*this);
     ASTContext &ASTContext = Actions.getASTContext();
 
     BalancedDelimiterTracker T(*this, tok::l_paren,
                                tok::annot_pragma_openmp_end);
     while (Tok.isNot(tok::annot_pragma_openmp_end)) {
+      TentativeParsingAction TPA(*this);
       OpenMPClauseKind CKind = Tok.isAnnotation()
                                    ? OMPC_unknown
                                    : getOpenMPClauseKind(PP.getSpelling(Tok));
       SourceLocation Loc = ConsumeToken();
 
       // Parse '('.
       if (T.expectAndConsume(diag::err_expected_lparen_after,
                              getOpenMPClauseName(CKind).data()))
         return Directive;
 
       OMPTraitInfo &TI = Actions.getASTContext().getNewOMPTraitInfo();
+      if (CKind == OMPC_when || CKind == OMPC_default) {
+        // If it is a "when" clause parse the context selectors for the trait
+        // info, a "default" clause will have empty trait info.
         if (CKind == OMPC_when) {
           // parse and get OMPTraitInfo to pass to the When clause
           parseOMPContextSelectors(Loc, TI);
           if (TI.Sets.size() == 0) {
             Diag(Tok, diag::err_omp_expected_context_selector) << "when clause";
             TPA.Commit();
             return Directive;
           }
 
           // Parse ':'
           if (Tok.is(tok::colon))
             ConsumeAnyToken();
           else {
             Diag(Tok, diag::err_omp_expected_colon) << "when clause";
             TPA.Commit();
             return Directive;
           }
-      }
-      // Skip Directive for now. We will parse directive in the second iteration
-      int paren = 0;
-      while (Tok.isNot(tok::r_paren) || paren != 0) {
-        if (Tok.is(tok::l_paren))
-          paren++;
-        if (Tok.is(tok::r_paren))
-          paren--;
-        if (Tok.is(tok::annot_pragma_openmp_end)) {
-          Diag(Tok, diag::err_omp_expected_punc)
-              << getOpenMPClauseName(CKind) << 0;
-          TPA.Commit();
-          return Directive;
-        }
-        ConsumeAnyToken();
-      }
-      // Parse ')'
-      if (Tok.is(tok::r_paren))
-        T.consumeClose();
 
           VariantMatchInfo VMI;
           TI.getAsVariantMatchInfo(ASTContext, VMI);
-
           VMIs.push_back(VMI);
         }
 
-    TPA.Revert();
-    // End of the first iteration. Parser is reset to the start of metadirective
-
-    std::function<void(StringRef)> DiagUnknownTrait = [this, Loc](
-                                                          StringRef ISATrait) {
-      // TODO Track the selector locations in a way that is accessible here to
-      // improve the diagnostic location.
-      Diag(Loc, diag::warn_unknown_declare_variant_isa_trait) << ISATrait;
-    };
-    TargetOMPContext OMPCtx(ASTContext, std::move(DiagUnknownTrait),
-                            /* CurrentFunctionDecl */ nullptr,
-                            ArrayRef<llvm::omp::TraitProperty>());
+        // TODO: currently expects a directive, OpenMP 5.1 specifies nothing
+        // directive when there is none.
+        ReadDirectiveWithinMetadirective = true;
+        StmtResult WhenDirective =
+            ParseOpenMPDeclarativeOrExecutableDirective(StmtCtx);
+        ReadDirectiveWithinMetadirective = false;
+        if (OMPExecutableDirective *D =
+                dyn_cast<OMPExecutableDirective>(WhenDirective.get()))
+          if (D->hasAssociatedStmt())
+            HasAssociatedStatement = true;
+
+        auto *WhenClause = Actions.ActOnOpenMPWhenClause(
+            TI, WhenDirective, Loc, T.getOpenLocation(), T.getCloseLocation());
+        Clauses.push_back(WhenClause);
 
-    // A single match is returned for OpenMP 5.0
-    int BestIdx = getBestVariantMatchForContext(VMIs, OMPCtx);
+        // Revert back to the beginning of the clause.
+        TPA.Revert(/*RevertAnnotations*/ true);
 
-    int Idx = 0;
-    // In OpenMP 5.0 metadirective is either replaced by another directive or
-    // ignored.
-    // TODO: In OpenMP 5.1 generate multiple directives based upon the matches
-    // found by getBestWhenMatchForContext.
-    while (Tok.isNot(tok::annot_pragma_openmp_end)) {
-      // OpenMP 5.0 implementation - Skip to the best index found.
-      if (Idx++ != BestIdx) {
+        // Skip until the end of this clause.
         ConsumeToken();  // Consume clause name
         T.consumeOpen(); // Consume '('
         int paren = 0;
-        // Skip everything inside the clause
         while (Tok.isNot(tok::r_paren) || paren != 0) {
           if (Tok.is(tok::l_paren))
             paren++;
           if (Tok.is(tok::r_paren))
             paren--;
+          if (Tok.is(tok::annot_pragma_openmp_end)) {
+            Diag(Tok, diag::err_omp_expected_punc)
+                << getOpenMPClauseName(CKind) << 0;
+            return Directive;
+          }
           ConsumeAnyToken();
         }
-        // Parse ')'
-        if (Tok.is(tok::r_paren))
-          T.consumeClose();
-        continue;
+
+        assert(Tok.is(tok::r_paren) &&
+               "Expected right paren ending the clause.");
       }
-
-      OpenMPClauseKind CKind = Tok.isAnnotation()
-                                   ? OMPC_unknown
-                                   : getOpenMPClauseKind(PP.getSpelling(Tok));
-      SourceLocation Loc = ConsumeToken();
-
-      // Parse '('.
-      T.consumeOpen();
-
-      // Skip ContextSelectors for when clause
-      if (CKind == OMPC_when) {
-        OMPTraitInfo &TI = Actions.getASTContext().getNewOMPTraitInfo();
-        // parse and skip the ContextSelectors
-        parseOMPContextSelectors(Loc, TI);
-
-        // Parse ':'
-        ConsumeAnyToken();
-      }
 
-      // If no directive is passed, skip in OpenMP 5.0.
-      // TODO: Generate nothing directive from OpenMP 5.1.
-      if (Tok.is(tok::r_paren)) {
-        SkipUntil(tok::annot_pragma_openmp_end);
-        break;
-      }
+    // Skip until the end of the metadirective.
+    // Skip any associated statement.
+    if (HasAssociatedStatement)
+      ParseStatement();
+
+    std::function<void(StringRef)> DiagUnknownTrait =
+        [this, Loc](StringRef ISATrait) {
+          // TODO Track the selector locations in a way that is accessible here
+          // to improve the diagnostic location.
+          Diag(Loc, diag::warn_unknown_declare_variant_isa_trait) << ISATrait;
+        };
+    TargetOMPContext OMPCtx(ASTContext, std::move(DiagUnknownTrait),
+                            /* CurrentFunctionDecl */ nullptr,
+                            ArrayRef<llvm::omp::TraitProperty>());
 
-      // Parse Directive
-      ReadDirectiveWithinMetadirective = true;
-      Directive = ParseOpenMPDeclarativeOrExecutableDirective(StmtCtx);
-      ReadDirectiveWithinMetadirective = false;
-      break;
+    // Find applicable clauses.
+    SmallVector<OMPClause *, 5> ApplicableClauses;
+    for (OMPClause *C : Clauses) {
+      OMPTraitInfo &TI = cast<OMPWhenClause>(C)->getTraitInfo();
+      VariantMatchInfo VMI;
+      TI.getAsVariantMatchInfo(ASTContext, VMI);
+      SmallVector<unsigned, 8> ConstructMatches;
+      if (isVariantApplicableInContext(VMI, OMPCtx, /*DeviceOnly*/ false))
+        ApplicableClauses.push_back(C);
     }
-    break;
+
+    return Actions.ActOnOpenMPMetaDirective(ApplicableClauses, Loc,
+                                            Tok.getLastLoc());
   }
   case OMPD_threadprivate: {
     // FIXME: Should this be permitted in C++?
     if ((StmtCtx & ParsedStmtContext::AllowDeclarationsInC) ==
         ParsedStmtContext()) {
       Diag(Tok, diag::err_omp_immediate_directive)
           << getOpenMPDirectiveName(DKind) << 0;
     }

clang/lib/Sema/SemaOpenMP.cpp

Diag(KindLoc, diag::err_omp_unexpected_clause_value)
                                    /*Last=*/unsigned(OMPC_BIND_unknown))
         << getOpenMPClauseName(OMPC_bind);
     return nullptr;
   }
 
   return OMPBindClause::Create(Context, Kind, KindLoc, StartLoc, LParenLoc,
                                EndLoc);
 }
+
+OMPClause *Sema::ActOnOpenMPWhenClause(OMPTraitInfo &TI, StmtResult Directive,
+                                       SourceLocation StartLoc,
+                                       SourceLocation LParenLoc,
+                                       SourceLocation EndLoc) {
+  return new (Context)
+      OMPWhenClause(TI, Directive.get(), StartLoc, LParenLoc, EndLoc);
+}
+
+StmtResult Sema::ActOnOpenMPMetaDirective(ArrayRef<OMPClause *> Clauses,
+                                          SourceLocation StartLoc,
+                                          SourceLocation EndLoc) {
+  return OMPMetaDirective::Create(Context, StartLoc, EndLoc, Clauses);
+}

clang/lib/Sema/TreeTransform.h

   /// By default, performs semantic analysis to build the new OpenMP clause.
   /// Subclasses may override this routine to provide different behavior.
   OMPClause *RebuildOMPAlignClause(Expr *A, SourceLocation StartLoc,
                                    SourceLocation LParenLoc,
                                    SourceLocation EndLoc) {
     return getSema().ActOnOpenMPAlignClause(A, StartLoc, LParenLoc, EndLoc);
   }
 
+  /// Build a new OpenMP 'when' clause.
+  ///
+  /// By default, performs semantic analysis to build the new OpenMP clause.
+  /// Subclasses may override this routine to provide different behavior.
+  OMPClause *RebuildOMPWhenClause(OMPTraitInfo &TI, Stmt *Directive,
+                                  SourceLocation StartLoc,
+                                  SourceLocation LParenLoc,
+                                  SourceLocation EndLoc) {
+    return getSema().ActOnOpenMPWhenClause(TI, Directive, StartLoc, LParenLoc,
+                                           EndLoc);
+  }
+
   /// Rebuild the operand to an Objective-C \@synchronized statement.
   ///
   /// By default, performs semantic analysis to build the new statement.
   /// Subclasses may override this routine to provide different behavior.
   ExprResult RebuildObjCAtSynchronizedOperand(SourceLocation atLoc,
                                               Expr *object) {
     return getSema().ActOnObjCAtSynchronizedOperand(atLoc, object);
   }

 
 template <typename Derived>
 OMPClause *TreeTransform<Derived>::TransformOMPBindClause(OMPBindClause *C) {
   return getDerived().RebuildOMPBindClause(
       C->getBindKind(), C->getBindKindLoc(), C->getBeginLoc(),
       C->getLParenLoc(), C->getEndLoc());
 }
 
+template <typename Derived>
+OMPClause *TreeTransform<Derived>::TransformOMPWhenClause(OMPWhenClause *C) {
+  return getDerived().RebuildOMPWhenClause(C->getTraitInfo(), C->getDirective(),
+                                           C->getBeginLoc(), C->getLParenLoc(),
+                                           C->getEndLoc());
+}
+
 //===----------------------------------------------------------------------===//
 // Expression transformation
 //===----------------------------------------------------------------------===//
 template<typename Derived>
 ExprResult
 TreeTransform<Derived>::TransformConstantExpr(ConstantExpr *E) {
   return TransformExpr(E->getSubExpr());
 }

clang/lib/Serialization/ASTReader.cpp

 }
 
 void OMPClauseReader::VisitOMPBindClause(OMPBindClause *C) {
   C->setBindKind(Record.readEnum<OpenMPBindClauseKind>());
   C->setLParenLoc(Record.readSourceLocation());
   C->setBindKindLoc(Record.readSourceLocation());
 }
 
+void OMPClauseReader::VisitOMPWhenClause(OMPWhenClause *C) {
+  // TODO: check, not familiar with this.
+  C->setTraitInfo(Record.readOMPTraitInfo());
+  C->setDirective(Record.readStmt());
+}
+
 void OMPClauseReader::VisitOMPAlignClause(OMPAlignClause *C) {
   C->setAlignment(Record.readExpr());
   C->setLParenLoc(Record.readSourceLocation());
 }
 
 OMPTraitInfo *ASTRecordReader::readOMPTraitInfo() {
   OMPTraitInfo &TI = getContext().getNewOMPTraitInfo();
   TI.Sets.resize(readUInt32());

clang/lib/Serialization/ASTWriter.cpp

 }
 
 void OMPClauseWriter::VisitOMPBindClause(OMPBindClause *C) {
   Record.writeEnum(C->getBindKind());
   Record.AddSourceLocation(C->getLParenLoc());
   Record.AddSourceLocation(C->getBindKindLoc());
 }
 
+void OMPClauseWriter::VisitOMPWhenClause(OMPWhenClause *C) {
+  // TODO: check, not familiar with this.
+  Record.writeOMPTraitInfo(&C->getTraitInfo());
+  Record.AddStmt(C->getDirective());
+}
+
 void ASTRecordWriter::writeOMPTraitInfo(const OMPTraitInfo *TI) {
   writeUInt32(TI->Sets.size());
   for (const auto &Set : TI->Sets) {
     writeEnum(Set.Kind);
     writeUInt32(Set.Selectors.size());
     for (const auto &Selector : Set.Selectors) {
       writeEnum(Selector.Kind);
       writeBool(Selector.ScoreOrCondition);

clang/tools/libclang/CIndex.cpp

 }
 void OMPClauseEnqueue::VisitOMPAffinityClause(const OMPAffinityClause *C) {
   Visitor->AddStmt(C->getModifier());
   for (const Expr *E : C->varlists())
     Visitor->AddStmt(E);
 }
 void OMPClauseEnqueue::VisitOMPBindClause(const OMPBindClause *C) {}
 
+void OMPClauseEnqueue::VisitOMPWhenClause(const OMPWhenClause *C) {
+  Visitor->AddStmt(C->getDirective());
+}
+
 } // namespace
 
 void EnqueueVisitor::EnqueueChildren(const OMPClause *S) {
   unsigned size = WL.size();
   OMPClauseEnqueue Visitor(this);
   Visitor.Visit(S);
   if (size == WL.size())
     return;

llvm/include/llvm/Frontend/OpenMP/OMP.td

 def OMPC_Notinbranch : Clause<"notinbranch"> {}
 def OMPC_Filter : Clause<"filter"> {
   let clangClass = "OMPFilterClause";
   let flangClass = "ScalarIntExpr";
 }
 def OMPC_Align : Clause<"align"> {
   let clangClass = "OMPAlignClause";
 }
-def OMPC_When: Clause<"when"> {}
+def OMPC_When: Clause<"when"> {
+  let clangClass = "OMPWhenClause";
+}
 
 def OMPC_Bind : Clause<"bind"> {
   let clangClass = "OMPBindClause";
 }
 
 //===----------------------------------------------------------------------===//
 // Definition of OpenMP directives
 //===----------------------------------------------------------------------===//