[OpenACC] Loop construct basic Sema and AST work (#93742)

This patch implements the 'loop' construct AST, as well as the basic
appertainment rule. Additionally, it sets up the 'parent' compute
construct, which is necessary for codegen/other diagnostics.

A 'loop' can apply to a for or range-for loop, otherwise it has no other
restrictions (though some of its clauses do).

Author: erichkeane

clang/include/clang-c/Index.h

@@ -2150,7 +2150,11 @@ enum CXCursorKind {
    */
   CXCursor_OpenACCComputeConstruct = 320,
 
-  CXCursor_LastStmt = CXCursor_OpenACCComputeConstruct,
+  /** OpenACC Loop Construct.
+   */
+  CXCursor_OpenACCLoopConstruct = 321,
+
+  CXCursor_LastStmt = CXCursor_OpenACCLoopConstruct,
 
   /**
    * Cursor that represents the translation unit itself.

clang/include/clang/AST/RecursiveASTVisitor.h

@@ -4000,6 +4000,8 @@ bool RecursiveASTVisitor<Derived>::VisitOpenACCClauseList(
 
 DEF_TRAVERSE_STMT(OpenACCComputeConstruct,
                   { TRY_TO(TraverseOpenACCAssociatedStmtConstruct(S)); })
+DEF_TRAVERSE_STMT(OpenACCLoopConstruct,
+                  { TRY_TO(TraverseOpenACCAssociatedStmtConstruct(S)); })
 
 // FIXME: look at the following tricky-seeming exprs to see if we
 // need to recurse on anything.  These are ones that have methods

clang/include/clang/AST/StmtOpenACC.h

@@ -113,6 +113,8 @@ class OpenACCAssociatedStmtConstruct : public OpenACCConstructStmt {
     return const_cast<OpenACCAssociatedStmtConstruct *>(this)->children();
   }
 };
+
+class OpenACCLoopConstruct;
 /// This class represents a compute construct, representing a 'Kind' of
 /// `parallel', 'serial', or 'kernel'. These constructs are associated with a
 /// 'structured block', defined as:
@@ -165,6 +167,11 @@ class OpenACCComputeConstruct final
   }
 
   void setStructuredBlock(Stmt *S) { setAssociatedStmt(S); }
+  // Serialization helper function that searches the structured block for 'loop'
+  // constructs that should be associated with this, and sets their parent
+  // compute construct to this one. This isn't necessary normally, since we have
+  // the ability to record the state during parsing.
+  void findAndSetChildLoops();
 
 public:
   static bool classof(const Stmt *T) {
@@ -176,12 +183,74 @@ class OpenACCComputeConstruct final
   static OpenACCComputeConstruct *
   Create(const ASTContext &C, OpenACCDirectiveKind K, SourceLocation BeginLoc,
          SourceLocation DirectiveLoc, SourceLocation EndLoc,
-         ArrayRef<const OpenACCClause *> Clauses, Stmt *StructuredBlock);
+         ArrayRef<const OpenACCClause *> Clauses, Stmt *StructuredBlock,
+         ArrayRef<OpenACCLoopConstruct *> AssociatedLoopConstructs);
 
   Stmt *getStructuredBlock() { return getAssociatedStmt(); }
   const Stmt *getStructuredBlock() const {
     return const_cast<OpenACCComputeConstruct *>(this)->getStructuredBlock();
   }
 };
+/// This class represents a 'loop' construct.  The 'loop' construct applies to a
+/// 'for' loop (or range-for loop), and is optionally associated with a Compute
+/// Construct.
+class OpenACCLoopConstruct final
+    : public OpenACCAssociatedStmtConstruct,
+      public llvm::TrailingObjects<OpenACCLoopConstruct,
+                                   const OpenACCClause *> {
+  // The compute construct this loop is associated with, or nullptr if this is
+  // an orphaned loop construct, or if it hasn't been set yet.  Because we
+  // construct the directives at the end of their statement, the 'parent'
+  // construct is not yet available at the time of construction, so this needs
+  // to be set 'later'.
+  const OpenACCComputeConstruct *ParentComputeConstruct = nullptr;
+
+  friend class ASTStmtWriter;
+  friend class ASTStmtReader;
+  friend class ASTContext;
+  friend class OpenACCComputeConstruct;
+
+  OpenACCLoopConstruct(unsigned NumClauses);
+
+  OpenACCLoopConstruct(SourceLocation Start, SourceLocation DirLoc,
+                       SourceLocation End,
+                       ArrayRef<const OpenACCClause *> Clauses, Stmt *Loop);
+  void setLoop(Stmt *Loop);
+
+  void setParentComputeConstruct(OpenACCComputeConstruct *CC) {
+    assert(!ParentComputeConstruct && "Parent already set?");
+    ParentComputeConstruct = CC;
+  }
+
+public:
+  static bool classof(const Stmt *T) {
+    return T->getStmtClass() == OpenACCLoopConstructClass;
+  }
+
+  static OpenACCLoopConstruct *CreateEmpty(const ASTContext &C,
+                                           unsigned NumClauses);
+
+  static OpenACCLoopConstruct *
+  Create(const ASTContext &C, SourceLocation BeginLoc, SourceLocation DirLoc,
+         SourceLocation EndLoc, ArrayRef<const OpenACCClause *> Clauses,
+         Stmt *Loop);
+
+  Stmt *getLoop() { return getAssociatedStmt(); }
+  const Stmt *getLoop() const {
+    return const_cast<OpenACCLoopConstruct *>(this)->getLoop();
+  }
+
+  /// OpenACC 3.3 2.9:
+  /// An orphaned loop construct is a loop construct that is not lexically
+  /// enclosed within a compute construct. The parent compute construct of a
+  /// loop construct is the nearest compute construct that lexically contains
+  /// the loop construct.
+  bool isOrphanedLoopConstruct() const {
+    return ParentComputeConstruct == nullptr;
+  }
+  const OpenACCComputeConstruct *getParentComputeConstruct() const {
+    return ParentComputeConstruct;
+  }
+};
 } // namespace clang
 #endif // LLVM_CLANG_AST_STMTOPENACC_H

clang/include/clang/AST/TextNodeDumper.h

@@ -408,6 +408,7 @@ class TextNodeDumper
   VisitLifetimeExtendedTemporaryDecl(const LifetimeExtendedTemporaryDecl *D);
   void VisitHLSLBufferDecl(const HLSLBufferDecl *D);
   void VisitOpenACCConstructStmt(const OpenACCConstructStmt *S);
+  void VisitOpenACCLoopConstruct(const OpenACCLoopConstruct *S);
 };
 
 } // namespace clang

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -12413,6 +12413,9 @@ def err_acc_reduction_composite_type
 def err_acc_reduction_composite_member_type :Error<
     "OpenACC 'reduction' composite variable must not have non-scalar field">;
 def note_acc_reduction_composite_member_loc : Note<"invalid field is here">;
+def err_acc_loop_not_for_loop
+    : Error<"OpenACC 'loop' construct can only be applied to a 'for' loop">;
+def note_acc_construct_here : Note<"'%0' construct is here">;
 
 // AMDGCN builtins diagnostics
 def err_amdgcn_global_load_lds_size_invalid_value : Error<"invalid size value">;

clang/include/clang/Basic/StmtNodes.td

@@ -302,3 +302,4 @@ def OpenACCConstructStmt : StmtNode<Stmt, /*abstract=*/1>;
 def OpenACCAssociatedStmtConstruct
     : StmtNode<OpenACCConstructStmt, /*abstract=*/1>;
 def OpenACCComputeConstruct : StmtNode<OpenACCAssociatedStmtConstruct>;
+def OpenACCLoopConstruct : StmtNode<OpenACCAssociatedStmtConstruct>;

clang/include/clang/Sema/SemaOpenACC.h

@@ -15,6 +15,7 @@
 #define LLVM_CLANG_SEMA_SEMAOPENACC_H
 
 #include "clang/AST/DeclGroup.h"
+#include "clang/AST/StmtOpenACC.h"
 #include "clang/Basic/OpenACCKinds.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Sema/Ownership.h"
@@ -25,6 +26,15 @@ namespace clang {
 class OpenACCClause;
 
 class SemaOpenACC : public SemaBase {
+private:
+  /// A collection of loop constructs in the compute construct scope that
+  /// haven't had their 'parent' compute construct set yet. Entires will only be
+  /// made to this list in the case where we know the loop isn't an orphan.
+  llvm::SmallVector<OpenACCLoopConstruct *> ParentlessLoopConstructs;
+  /// Whether we are inside of a compute construct, and should add loops to the
+  /// above collection.
+  bool InsideComputeConstruct = false;
+
 public:
   // Redeclaration of the version in OpenACCClause.h.
   using DeviceTypeArgument = std::pair<IdentifierInfo *, SourceLocation>;
@@ -394,7 +404,8 @@ class SemaOpenACC : public SemaBase {
   bool ActOnStartDeclDirective(OpenACCDirectiveKind K, SourceLocation StartLoc);
   /// Called when we encounter an associated statement for our construct, this
   /// should check legality of the statement as it appertains to this Construct.
-  StmtResult ActOnAssociatedStmt(OpenACCDirectiveKind K, StmtResult AssocStmt);
+  StmtResult ActOnAssociatedStmt(SourceLocation DirectiveLoc,
+                                 OpenACCDirectiveKind K, StmtResult AssocStmt);
 
   /// Called after the directive has been completely parsed, including the
   /// declaration group or associated statement.
@@ -431,6 +442,20 @@ class SemaOpenACC : public SemaBase {
                                    Expr *LowerBound,
                                    SourceLocation ColonLocFirst, Expr *Length,
                                    SourceLocation RBLoc);
+
+  /// 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.
+  class AssociatedStmtRAII {
+    SemaOpenACC &SemaRef;
+    bool WasInsideComputeConstruct;
+    OpenACCDirectiveKind DirKind;
+    llvm::SmallVector<OpenACCLoopConstruct *> ParentlessLoopConstructs;
+
+  public:
+    AssociatedStmtRAII(SemaOpenACC &, OpenACCDirectiveKind);
+    ~AssociatedStmtRAII();
+  };
 };
 
 } // namespace clang

clang/include/clang/Serialization/ASTBitCodes.h

@@ -1946,6 +1946,7 @@ enum StmtCode {
 
   // OpenACC Constructs
   STMT_OPENACC_COMPUTE_CONSTRUCT,
+  STMT_OPENACC_LOOP_CONSTRUCT,
 };
 
 /// The kinds of designators that can occur in a

clang/lib/AST/StmtOpenACC.cpp

@@ -12,6 +12,8 @@
 
 #include "clang/AST/StmtOpenACC.h"
 #include "clang/AST/ASTContext.h"
+#include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/AST/StmtCXX.h"
 using namespace clang;
 
 OpenACCComputeConstruct *
@@ -26,11 +28,98 @@ OpenACCComputeConstruct::CreateEmpty(const ASTContext &C, unsigned NumClauses) {
 OpenACCComputeConstruct *OpenACCComputeConstruct::Create(
     const ASTContext &C, OpenACCDirectiveKind K, SourceLocation BeginLoc,
     SourceLocation DirLoc, SourceLocation EndLoc,
-    ArrayRef<const OpenACCClause *> Clauses, Stmt *StructuredBlock) {
+    ArrayRef<const OpenACCClause *> Clauses, Stmt *StructuredBlock,
+    ArrayRef<OpenACCLoopConstruct *> AssociatedLoopConstructs) {
   void *Mem = C.Allocate(
       OpenACCComputeConstruct::totalSizeToAlloc<const OpenACCClause *>(
           Clauses.size()));
   auto *Inst = new (Mem) OpenACCComputeConstruct(K, BeginLoc, DirLoc, EndLoc,
                                                  Clauses, StructuredBlock);
+
+  llvm::for_each(AssociatedLoopConstructs, [&](OpenACCLoopConstruct *C) {
+    C->setParentComputeConstruct(Inst);
+  });
+
+  return Inst;
+}
+
+void OpenACCComputeConstruct::findAndSetChildLoops() {
+  struct LoopConstructFinder : RecursiveASTVisitor<LoopConstructFinder> {
+    OpenACCComputeConstruct *Construct = nullptr;
+
+    LoopConstructFinder(OpenACCComputeConstruct *Construct)
+        : Construct(Construct) {}
+
+    bool TraverseOpenACCComputeConstruct(OpenACCComputeConstruct *C) {
+      // Stop searching if we find a compute construct.
+      return true;
+    }
+    bool TraverseOpenACCLoopConstruct(OpenACCLoopConstruct *C) {
+      // Stop searching if we find a loop construct, after taking ownership of
+      // it.
+      C->setParentComputeConstruct(Construct);
+      return true;
+    }
+  };
+
+  LoopConstructFinder f(this);
+  f.TraverseStmt(getAssociatedStmt());
+}
+
+OpenACCLoopConstruct::OpenACCLoopConstruct(unsigned NumClauses)
+    : OpenACCAssociatedStmtConstruct(
+          OpenACCLoopConstructClass, OpenACCDirectiveKind::Loop,
+          SourceLocation{}, SourceLocation{}, SourceLocation{},
+          /*AssociatedStmt=*/nullptr) {
+  std::uninitialized_value_construct(
+      getTrailingObjects<const OpenACCClause *>(),
+      getTrailingObjects<const OpenACCClause *>() + NumClauses);
+  setClauseList(
+      MutableArrayRef(getTrailingObjects<const OpenACCClause *>(), NumClauses));
+}
+
+OpenACCLoopConstruct::OpenACCLoopConstruct(
+    SourceLocation Start, SourceLocation DirLoc, SourceLocation End,
+    ArrayRef<const OpenACCClause *> Clauses, Stmt *Loop)
+    : OpenACCAssociatedStmtConstruct(OpenACCLoopConstructClass,
+                                     OpenACCDirectiveKind::Loop, Start, DirLoc,
+                                     End, Loop) {
+  // accept 'nullptr' for the loop. This is diagnosed somewhere, but this gives
+  // us some level of AST fidelity in the error case.
+  assert((Loop == nullptr || isa<ForStmt, CXXForRangeStmt>(Loop)) &&
+         "Associated Loop not a for loop?");
+  // Initialize the trailing storage.
+  std::uninitialized_copy(Clauses.begin(), Clauses.end(),
+                          getTrailingObjects<const OpenACCClause *>());
+
+  setClauseList(MutableArrayRef(getTrailingObjects<const OpenACCClause *>(),
+                                Clauses.size()));
+}
+
+void OpenACCLoopConstruct::setLoop(Stmt *Loop) {
+  assert((isa<ForStmt, CXXForRangeStmt>(Loop)) &&
+         "Associated Loop not a for loop?");
+  setAssociatedStmt(Loop);
+}
+
+OpenACCLoopConstruct *OpenACCLoopConstruct::CreateEmpty(const ASTContext &C,
+                                                        unsigned NumClauses) {
+  void *Mem =
+      C.Allocate(OpenACCLoopConstruct::totalSizeToAlloc<const OpenACCClause *>(
+          NumClauses));
+  auto *Inst = new (Mem) OpenACCLoopConstruct(NumClauses);
+  return Inst;
+}
+
+OpenACCLoopConstruct *
+OpenACCLoopConstruct::Create(const ASTContext &C, SourceLocation BeginLoc,
+                             SourceLocation DirLoc, SourceLocation EndLoc,
+                             ArrayRef<const OpenACCClause *> Clauses,
+                             Stmt *Loop) {
+  void *Mem =
+      C.Allocate(OpenACCLoopConstruct::totalSizeToAlloc<const OpenACCClause *>(
+          Clauses.size()));
+  auto *Inst =
+      new (Mem) OpenACCLoopConstruct(BeginLoc, DirLoc, EndLoc, Clauses, Loop);
   return Inst;
 }

clang/lib/AST/StmtPrinter.cpp

@@ -1156,6 +1156,19 @@ void StmtPrinter::VisitOpenACCComputeConstruct(OpenACCComputeConstruct *S) {
   PrintStmt(S->getStructuredBlock());
 }
 
+void StmtPrinter::VisitOpenACCLoopConstruct(OpenACCLoopConstruct *S) {
+  Indent() << "#pragma acc loop";
+
+  if (!S->clauses().empty()) {
+    OS << ' ';
+    OpenACCClausePrinter Printer(OS, Policy);
+    Printer.VisitClauseList(S->clauses());
+  }
+  OS << '\n';
+
+  PrintStmt(S->getLoop());
+}
+
 //===----------------------------------------------------------------------===//
 //  Expr printing methods.
 //===----------------------------------------------------------------------===//