[OpenACC] Enforce all 'collapse' intervening rules (llvm#110684)

'collapse' makes the 'loop' construct apply to the next N nested loops,
and 'loop' requires all associated loops be 'for' loops (or range-for).

This patch adds this restriction, plus adds a number of infrastructure
changes to permit some of the other restrictions in the future to be
implemented.

'collapse' also requires that there is not any calls to other directives
 inside of the collapse region, so this also implements that limitation.

Author: erichkeane

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -12617,6 +12617,11 @@ def err_acc_loop_spec_conflict
 def err_acc_collapse_loop_count
     : Error<"OpenACC 'collapse' clause loop count must be a %select{constant "
             "expression|positive integer value, evaluated to %1}0">;
+def err_acc_invalid_in_collapse_loop
+    : Error<"%select{OpenACC '%1' construct|while loop|do loop}0 cannot appear "
+            "in intervening code of a 'loop' with a 'collapse' clause">;
+def note_acc_collapse_clause_here
+    : Note<"active 'collapse' clause defined here">;
 
 // AMDGCN builtins diagnostics
 def err_amdgcn_global_load_lds_size_invalid_value : Error<"invalid size value">;

clang/include/clang/Sema/SemaOpenACC.h

@@ -22,6 +22,7 @@
 #include "clang/Sema/Ownership.h"
 #include "clang/Sema/SemaBase.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Compiler.h"
 #include <cassert>
 #include <optional>
 #include <utility>
@@ -41,6 +42,30 @@ class SemaOpenACC : public SemaBase {
   /// above collection.
   bool InsideComputeConstruct = false;
 
+  /// The 'collapse' clause requires quite a bit of checking while
+  /// parsing/instantiating its body, so this structure/object keeps all of the
+  /// necessary information as we do checking.  This should rarely be directly
+  /// modified, and typically should be controlled by the RAII objects.
+  ///
+  /// Collapse has an 'N' count that makes it apply to a number of loops 'below'
+  /// it.
+  struct CollapseCheckingInfo {
+    OpenACCCollapseClause *ActiveCollapse = nullptr;
+
+    /// This is a value that maintains the current value of the 'N' on the
+    /// current collapse, minus the depth that has already been traversed. When
+    /// there is not an active collapse, or a collapse whose depth we don't know
+    /// (for example, if it is a dependent value), this should be `nullopt`,
+    /// else it should be 'N' minus the current depth traversed.
+    std::optional<llvm::APSInt> CurCollapseCount;
+
+    /// Records whether we've seen the top level 'for'. We already diagnose
+    /// later that the 'top level' is a for loop, so we use this to suppress the
+    /// 'collapse inner loop not a 'for' loop' diagnostic.
+    LLVM_PREFERRED_TYPE(bool)
+    unsigned TopLevelLoopSeen : 1;
+  } CollapseInfo{nullptr, std::nullopt, false};
+
 public:
   // Redeclaration of the version in OpenACCClause.h.
   using DeviceTypeArgument = std::pair<IdentifierInfo *, SourceLocation>;
@@ -411,6 +436,17 @@ class SemaOpenACC : public SemaBase {
 
   SemaOpenACC(Sema &S);
 
+  // Called when we encounter a 'while' statement, before looking at its 'body'.
+  void ActOnWhileStmt(SourceLocation WhileLoc);
+  // Called when we encounter a 'do' statement, before looking at its 'body'.
+  void ActOnDoStmt(SourceLocation DoLoc);
+  // Called when we encounter a 'for' statement, before looking at its 'body'.
+  void ActOnForStmtBegin(SourceLocation ForLoc);
+  // Called when we encounter a 'for' statement, after we've consumed/checked
+  // the body. This is necessary for a number of checks on the contents of the
+  // 'for' statement.
+  void ActOnForStmtEnd(SourceLocation ForLoc, StmtResult Body);
+
   /// Called after parsing an OpenACC Clause so that it can be checked.
   OpenACCClause *ActOnClause(ArrayRef<const OpenACCClause *> ExistingClauses,
                              OpenACCParsedClause &Clause);
@@ -474,6 +510,18 @@ class SemaOpenACC : public SemaBase {
   /// Checks the loop depth value for a collapse clause.
   ExprResult CheckCollapseLoopCount(Expr *LoopCount);
 
+  /// Helper type to restore the state of various 'loop' constructs when we run
+  /// into a loop (for, etc) inside the construct.
+  class LoopInConstructRAII {
+    SemaOpenACC &SemaRef;
+    CollapseCheckingInfo OldCollapseInfo;
+
+  public:
+    LoopInConstructRAII(SemaOpenACC &SemaRef)
+        : SemaRef(SemaRef), OldCollapseInfo(SemaRef.CollapseInfo) {}
+    ~LoopInConstructRAII() { SemaRef.CollapseInfo = OldCollapseInfo; }
+  };
+
   /// Helper type for the registration/assignment of constructs that need to
   /// 'know' about their parent constructs and hold a reference to them, such as
   /// Loop needing its parent construct.
@@ -482,9 +530,15 @@ class SemaOpenACC : public SemaBase {
     bool WasInsideComputeConstruct;
     OpenACCDirectiveKind DirKind;
     llvm::SmallVector<OpenACCLoopConstruct *> ParentlessLoopConstructs;
+    LoopInConstructRAII LoopRAII;
 
   public:
-    AssociatedStmtRAII(SemaOpenACC &, OpenACCDirectiveKind);
+    AssociatedStmtRAII(SemaOpenACC &, OpenACCDirectiveKind,
+                       ArrayRef<const OpenACCClause *>,
+                       ArrayRef<OpenACCClause *>);
+    void SetCollapseInfoBeforeAssociatedStmt(
+        ArrayRef<const OpenACCClause *> UnInstClauses,
+        ArrayRef<OpenACCClause *> Clauses);
     ~AssociatedStmtRAII();
   };
 };

clang/lib/Parse/ParseOpenACC.cpp

@@ -1461,8 +1461,8 @@ StmtResult Parser::ParseOpenACCDirectiveStmt() {
     return StmtError();
 
   StmtResult AssocStmt;
-  SemaOpenACC::AssociatedStmtRAII AssocStmtRAII(getActions().OpenACC(),
-                                                DirInfo.DirKind);
+  SemaOpenACC::AssociatedStmtRAII AssocStmtRAII(
+      getActions().OpenACC(), DirInfo.DirKind, {}, DirInfo.Clauses);
   if (doesDirectiveHaveAssociatedStmt(DirInfo.DirKind)) {
     ParsingOpenACCDirectiveRAII DirScope(*this, /*Value=*/false);
     ParseScope ACCScope(this, getOpenACCScopeFlags(DirInfo.DirKind));

clang/lib/Parse/ParseStmt.cpp

@@ -24,6 +24,7 @@
 #include "clang/Sema/Scope.h"
 #include "clang/Sema/SemaCodeCompletion.h"
 #include "clang/Sema/SemaObjC.h"
+#include "clang/Sema/SemaOpenACC.h"
 #include "clang/Sema/SemaOpenMP.h"
 #include "clang/Sema/TypoCorrection.h"
 #include "llvm/ADT/STLExtras.h"
@@ -1854,6 +1855,11 @@ StmtResult Parser::ParseWhileStatement(SourceLocation *TrailingElseLoc) {
                                 Sema::ConditionKind::Boolean, LParen, RParen))
     return StmtError();
 
+  // OpenACC Restricts a while-loop inside of certain construct/clause
+  // combinations, so diagnose that here in OpenACC mode.
+  SemaOpenACC::LoopInConstructRAII LCR{getActions().OpenACC()};
+  getActions().OpenACC().ActOnWhileStmt(WhileLoc);
+
   // C99 6.8.5p5 - In C99, the body of the while statement is a scope, even if
   // there is no compound stmt.  C90 does not have this clause.  We only do this
   // if the body isn't a compound statement to avoid push/pop in common cases.
@@ -1902,6 +1908,11 @@ StmtResult Parser::ParseDoStatement() {
 
   ParseScope DoScope(this, ScopeFlags);
 
+  // OpenACC Restricts a do-while-loop inside of certain construct/clause
+  // combinations, so diagnose that here in OpenACC mode.
+  SemaOpenACC::LoopInConstructRAII LCR{getActions().OpenACC()};
+  getActions().OpenACC().ActOnDoStmt(DoLoc);
+
   // C99 6.8.5p5 - In C99, the body of the do statement is a scope, even if
   // there is no compound stmt.  C90 does not have this clause. We only do this
   // if the body isn't a compound statement to avoid push/pop in common cases.
@@ -2326,6 +2337,11 @@ StmtResult Parser::ParseForStatement(SourceLocation *TrailingElseLoc) {
     }
   }
 
+  // OpenACC Restricts a for-loop inside of certain construct/clause
+  // combinations, so diagnose that here in OpenACC mode.
+  SemaOpenACC::LoopInConstructRAII LCR{getActions().OpenACC()};
+  getActions().OpenACC().ActOnForStmtBegin(ForLoc);
+
   // C99 6.8.5p5 - In C99, the body of the for statement is a scope, even if
   // there is no compound stmt.  C90 does not have this clause.  We only do this
   // if the body isn't a compound statement to avoid push/pop in common cases.
@@ -2358,6 +2374,8 @@ StmtResult Parser::ParseForStatement(SourceLocation *TrailingElseLoc) {
   // Pop the body scope if needed.
   InnerScope.Exit();
 
+  getActions().OpenACC().ActOnForStmtEnd(ForLoc, Body);
+
   // Leave the for-scope.
   ForScope.Exit();
 

clang/lib/Sema/SemaOpenACC.cpp

@@ -1073,19 +1073,62 @@ OpenACCClause *SemaOpenACCClauseVisitor::VisitCollapseClause(
 
 SemaOpenACC::SemaOpenACC(Sema &S) : SemaBase(S) {}
 
-SemaOpenACC::AssociatedStmtRAII::AssociatedStmtRAII(SemaOpenACC &S,
-                                                    OpenACCDirectiveKind DK)
+SemaOpenACC::AssociatedStmtRAII::AssociatedStmtRAII(
+    SemaOpenACC &S, OpenACCDirectiveKind DK,
+    ArrayRef<const OpenACCClause *> UnInstClauses,
+    ArrayRef<OpenACCClause *> Clauses)
     : SemaRef(S), WasInsideComputeConstruct(S.InsideComputeConstruct),
-      DirKind(DK) {
+      DirKind(DK), LoopRAII(SemaRef) {
   // Compute constructs end up taking their 'loop'.
   if (DirKind == OpenACCDirectiveKind::Parallel ||
       DirKind == OpenACCDirectiveKind::Serial ||
       DirKind == OpenACCDirectiveKind::Kernels) {
     SemaRef.InsideComputeConstruct = true;
     SemaRef.ParentlessLoopConstructs.swap(ParentlessLoopConstructs);
+  } else if (DirKind == OpenACCDirectiveKind::Loop) {
+    SetCollapseInfoBeforeAssociatedStmt(UnInstClauses, Clauses);
   }
 }
 
+void SemaOpenACC::AssociatedStmtRAII::SetCollapseInfoBeforeAssociatedStmt(
+    ArrayRef<const OpenACCClause *> UnInstClauses,
+    ArrayRef<OpenACCClause *> Clauses) {
+  // We make sure to take an optional list of uninstantiated clauses, so that
+  // we can check to make sure we don't 'double diagnose' in the event that
+  // the value of 'N' was not dependent in a template. We also ensure during
+  // Sema that there is only 1 collapse on each construct, so we can count on
+  // the fact that if both find a 'collapse', that they are the same one.
+  auto *CollapseClauseItr =
+      llvm::find_if(Clauses, llvm::IsaPred<OpenACCCollapseClause>);
+  auto *UnInstCollapseClauseItr =
+      llvm::find_if(UnInstClauses, llvm::IsaPred<OpenACCCollapseClause>);
+
+  if (Clauses.end() == CollapseClauseItr)
+    return;
+
+  OpenACCCollapseClause *CollapseClause =
+      cast<OpenACCCollapseClause>(*CollapseClauseItr);
+
+  SemaRef.CollapseInfo.ActiveCollapse = CollapseClause;
+  Expr *LoopCount = CollapseClause->getLoopCount();
+
+  // If the loop count is still instantiation dependent, setting the depth
+  // counter isn't necessary, so return here.
+  if (!LoopCount || LoopCount->isInstantiationDependent())
+    return;
+
+  // Suppress diagnostics if we've done a 'transform' where the previous version
+  // wasn't dependent, meaning we already diagnosed it.
+  if (UnInstCollapseClauseItr != UnInstClauses.end() &&
+      !cast<OpenACCCollapseClause>(*UnInstCollapseClauseItr)
+           ->getLoopCount()
+           ->isInstantiationDependent())
+    return;
+
+  SemaRef.CollapseInfo.CurCollapseCount =
+      cast<ConstantExpr>(LoopCount)->getResultAsAPSInt();
+}
+
 SemaOpenACC::AssociatedStmtRAII::~AssociatedStmtRAII() {
   SemaRef.InsideComputeConstruct = WasInsideComputeConstruct;
   if (DirKind == OpenACCDirectiveKind::Parallel ||
@@ -1094,6 +1137,9 @@ SemaOpenACC::AssociatedStmtRAII::~AssociatedStmtRAII() {
     assert(SemaRef.ParentlessLoopConstructs.empty() &&
            "Didn't consume loop construct list?");
     SemaRef.ParentlessLoopConstructs.swap(ParentlessLoopConstructs);
+  } else if (DirKind == OpenACCDirectiveKind::Loop) {
+    // Nothing really to do here, the LoopInConstruct should handle restorations
+    // correctly.
   }
 }
 
@@ -1646,10 +1692,78 @@ ExprResult SemaOpenACC::CheckCollapseLoopCount(Expr *LoopCount) {
       ConstantExpr::Create(getASTContext(), LoopCount, APValue{*ICE})};
 }
 
+void SemaOpenACC::ActOnWhileStmt(SourceLocation WhileLoc) {
+  if (!getLangOpts().OpenACC)
+    return;
+
+  if (!CollapseInfo.TopLevelLoopSeen)
+    return;
+
+  if (CollapseInfo.CurCollapseCount && *CollapseInfo.CurCollapseCount > 0) {
+    Diag(WhileLoc, diag::err_acc_invalid_in_collapse_loop) << /*while loop*/ 1;
+    assert(CollapseInfo.ActiveCollapse && "Collapse count without object?");
+    Diag(CollapseInfo.ActiveCollapse->getBeginLoc(),
+         diag::note_acc_collapse_clause_here);
+
+    // Remove the value so that we don't get cascading errors in the body. The
+    // caller RAII object will restore this.
+    CollapseInfo.CurCollapseCount = std::nullopt;
+  }
+}
+
+void SemaOpenACC::ActOnDoStmt(SourceLocation DoLoc) {
+  if (!getLangOpts().OpenACC)
+    return;
+
+  if (!CollapseInfo.TopLevelLoopSeen)
+    return;
+
+  if (CollapseInfo.CurCollapseCount && *CollapseInfo.CurCollapseCount > 0) {
+    Diag(DoLoc, diag::err_acc_invalid_in_collapse_loop) << /*do loop*/ 2;
+    assert(CollapseInfo.ActiveCollapse && "Collapse count without object?");
+    Diag(CollapseInfo.ActiveCollapse->getBeginLoc(),
+         diag::note_acc_collapse_clause_here);
+
+    // Remove the value so that we don't get cascading errors in the body. The
+    // caller RAII object will restore this.
+    CollapseInfo.CurCollapseCount = std::nullopt;
+  }
+}
+
+void SemaOpenACC::ActOnForStmtBegin(SourceLocation ForLoc) {
+  if (!getLangOpts().OpenACC)
+    return;
+
+  // Enable the while/do-while checking.
+  CollapseInfo.TopLevelLoopSeen = true;
+
+  if (CollapseInfo.CurCollapseCount && *CollapseInfo.CurCollapseCount > 0)
+    --(*CollapseInfo.CurCollapseCount);
+}
+
+void SemaOpenACC::ActOnForStmtEnd(SourceLocation ForLoc, StmtResult Body) {
+  if (!getLangOpts().OpenACC)
+    return;
+}
+
 bool SemaOpenACC::ActOnStartStmtDirective(OpenACCDirectiveKind K,
                                           SourceLocation StartLoc) {
   SemaRef.DiscardCleanupsInEvaluationContext();
   SemaRef.PopExpressionEvaluationContext();
+
+  // OpenACC 3.3 2.9.1:
+  // Intervening code must not contain other OpenACC directives or calls to API
+  // routines.
+  //
+  // ALL constructs are ill-formed if there is an active 'collapse'
+  if (CollapseInfo.CurCollapseCount && *CollapseInfo.CurCollapseCount > 0) {
+    Diag(StartLoc, diag::err_acc_invalid_in_collapse_loop)
+        << /*OpenACC Construct*/ 0 << K;
+    assert(CollapseInfo.ActiveCollapse && "Collapse count without object?");
+    Diag(CollapseInfo.ActiveCollapse->getBeginLoc(),
+         diag::note_acc_collapse_clause_here);
+  }
+
   return diagnoseConstructAppertainment(*this, K, StartLoc, /*IsStmt=*/true);
 }