https://github.com/loopacino created 
https://github.com/llvm/llvm-project/pull/206977

Adds the OpenMP 6.1 Section 17.4 `Flatten Construct`

>From 0fb5bd42f1e127faea963c359b1e3c69cb969c60 Mon Sep 17 00:00:00 2001
From: amtiwari <[email protected]>
Date: Wed, 1 Jul 2026 09:06:58 -0400
Subject: [PATCH] add OMPFlattenDirective ast node

---
 clang/include/clang/AST/RecursiveASTVisitor.h |  3 +
 clang/include/clang/AST/StmtOpenMP.h          | 75 ++++++++++++++++++-
 clang/include/clang/Basic/StmtNodes.td        |  1 +
 clang/lib/AST/StmtOpenMP.cpp                  | 20 +++++
 clang/lib/AST/StmtPrinter.cpp                 |  5 ++
 clang/lib/AST/StmtProfile.cpp                 |  4 +
 clang/lib/Basic/OpenMPKinds.cpp               |  2 +-
 7 files changed, 108 insertions(+), 2 deletions(-)

diff --git a/clang/include/clang/AST/RecursiveASTVisitor.h 
b/clang/include/clang/AST/RecursiveASTVisitor.h
index b000a34043696..e425f7a31f59c 100644
--- a/clang/include/clang/AST/RecursiveASTVisitor.h
+++ b/clang/include/clang/AST/RecursiveASTVisitor.h
@@ -3227,6 +3227,9 @@ DEF_TRAVERSE_STMT(OMPFuseDirective,
 DEF_TRAVERSE_STMT(OMPInterchangeDirective,
                   { TRY_TO(TraverseOMPExecutableDirective(S)); })
 
+DEF_TRAVERSE_STMT(OMPFlattenDirective,
+                  { TRY_TO(TraverseOMPExecutableDirective(S)); })
+
 DEF_TRAVERSE_STMT(OMPSplitDirective,
                   { TRY_TO(TraverseOMPExecutableDirective(S)); })
 
diff --git a/clang/include/clang/AST/StmtOpenMP.h 
b/clang/include/clang/AST/StmtOpenMP.h
index dbc76e7df8ecd..1420331f37c1f 100644
--- a/clang/include/clang/AST/StmtOpenMP.h
+++ b/clang/include/clang/AST/StmtOpenMP.h
@@ -1041,7 +1041,7 @@ class OMPCanonicalLoopNestTransformationDirective
     Stmt::StmtClass C = T->getStmtClass();
     return C == OMPTileDirectiveClass || C == OMPUnrollDirectiveClass ||
            C == OMPReverseDirectiveClass || C == OMPInterchangeDirectiveClass 
||
-           C == OMPStripeDirectiveClass;
+           C == OMPStripeDirectiveClass || C == OMPFlattenDirectiveClass;
   }
 };
 
@@ -5959,6 +5959,79 @@ class OMPInterchangeDirective final
   }
 };
 
+/// Represents the '#pragma omp flatten' loop transformation directive.
+///
+/// \code{c}
+///   #pragma omp flatten
+///   for (int i = 0; i < m; ++i)
+///     for (int j = 0; j < n; ++j)
+///       ..
+/// \endcode
+class OMPFlattenDirective final
+    : public OMPCanonicalLoopNestTransformationDirective {
+  friend class ASTStmtReader;
+  friend class OMPExecutableDirective;
+
+  /// Offsets of child members.
+  enum {
+    PreInitsOffset = 0,
+    TransformedStmtOffset,
+  };
+
+  explicit OMPFlattenDirective(SourceLocation StartLoc, SourceLocation EndLoc,
+                               unsigned NumLoops)
+      : OMPCanonicalLoopNestTransformationDirective(
+            OMPFlattenDirectiveClass, llvm::omp::OMPD_flatten, StartLoc, 
EndLoc,
+            NumLoops) {}
+
+  void setPreInits(Stmt *PreInits) {
+    Data->getChildren()[PreInitsOffset] = PreInits;
+  }
+
+  void setTransformedStmt(Stmt *S) {
+    Data->getChildren()[TransformedStmtOffset] = S;
+  }
+
+public:
+  /// Create a new AST node representation for '#pragma omp flatten'.
+  ///
+  /// \param C         Context of the AST.
+  /// \param StartLoc  Location of the introducer (e.g. the 'omp' token).
+  /// \param EndLoc    Location of the directive's end (e.g. the tok::eod).
+  /// \param Clauses   The directive's clauses.
+  /// \param NumLoops  Number of affected loops (the flatten depth, currently
+  ///                  always 2).
+  /// \param AssociatedStmt  The outermost associated loop.
+  /// \param TransformedStmt The flattened loop, or nullptr in dependent
+  ///                        contexts.
+  /// \param PreInits  Helper preinits statements for the loop nest.
+  static OMPFlattenDirective *
+  Create(const ASTContext &C, SourceLocation StartLoc, SourceLocation EndLoc,
+         ArrayRef<OMPClause *> Clauses, unsigned NumLoops, Stmt 
*AssociatedStmt,
+         Stmt *TransformedStmt, Stmt *PreInits);
+
+  /// Build an empty '#pragma omp flatten' AST node for deserialization.
+  ///
+  /// \param C          Context of the AST.
+  /// \param NumClauses Number of clauses to allocate.
+  /// \param NumLoops   Number of associated loops to allocate.
+  static OMPFlattenDirective *
+  CreateEmpty(const ASTContext &C, unsigned NumClauses, unsigned NumLoops);
+
+  /// Gets the flattened loop after the transformation. This is the de-sugared
+  /// replacement or nullptr in dependent contexts.
+  Stmt *getTransformedStmt() const {
+    return Data->getChildren()[TransformedStmtOffset];
+  }
+
+  /// Return preinits statement.
+  Stmt *getPreInits() const { return Data->getChildren()[PreInitsOffset]; }
+
+  static bool classof(const Stmt *T) {
+    return T->getStmtClass() == OMPFlattenDirectiveClass;
+  }
+};
+
 /// The base class for all transformation directives of canonical loop
 /// sequences (currently only 'fuse')
 class OMPCanonicalLoopSequenceTransformationDirective
diff --git a/clang/include/clang/Basic/StmtNodes.td 
b/clang/include/clang/Basic/StmtNodes.td
index e166894ea024b..0d7a0c1490a1e 100644
--- a/clang/include/clang/Basic/StmtNodes.td
+++ b/clang/include/clang/Basic/StmtNodes.td
@@ -247,6 +247,7 @@ def OMPReverseDirective : 
StmtNode<OMPCanonicalLoopNestTransformationDirective>;
 def OMPSplitDirective : StmtNode<OMPCanonicalLoopNestTransformationDirective>;
 def OMPInterchangeDirective
     : StmtNode<OMPCanonicalLoopNestTransformationDirective>;
+def OMPFlattenDirective : 
StmtNode<OMPCanonicalLoopNestTransformationDirective>;
 def OMPCanonicalLoopSequenceTransformationDirective
     : StmtNode<OMPExecutableDirective, 1>;
 def OMPFuseDirective
diff --git a/clang/lib/AST/StmtOpenMP.cpp b/clang/lib/AST/StmtOpenMP.cpp
index 9d6b315effb41..477604ef3dafb 100644
--- a/clang/lib/AST/StmtOpenMP.cpp
+++ b/clang/lib/AST/StmtOpenMP.cpp
@@ -552,6 +552,26 @@ OMPInterchangeDirective::CreateEmpty(const ASTContext &C, 
unsigned NumClauses,
       SourceLocation(), SourceLocation(), NumLoops);
 }
 
+OMPFlattenDirective *OMPFlattenDirective::Create(
+    const ASTContext &C, SourceLocation StartLoc, SourceLocation EndLoc,
+    ArrayRef<OMPClause *> Clauses, unsigned NumLoops, Stmt *AssociatedStmt,
+    Stmt *TransformedStmt, Stmt *PreInits) {
+  OMPFlattenDirective *Dir = createDirective<OMPFlattenDirective>(
+      C, Clauses, AssociatedStmt, TransformedStmtOffset + 1, StartLoc, EndLoc,
+      NumLoops);
+  Dir->setTransformedStmt(TransformedStmt);
+  Dir->setPreInits(PreInits);
+  return Dir;
+}
+
+OMPFlattenDirective *OMPFlattenDirective::CreateEmpty(const ASTContext &C,
+                                                      unsigned NumClauses,
+                                                      unsigned NumLoops) {
+  return createEmptyDirective<OMPFlattenDirective>(
+      C, NumClauses, /*HasAssociatedStmt=*/true, TransformedStmtOffset + 1,
+      SourceLocation(), SourceLocation(), NumLoops);
+}
+
 OMPSplitDirective *
 OMPSplitDirective::Create(const ASTContext &C, SourceLocation StartLoc,
                           SourceLocation EndLoc, ArrayRef<OMPClause *> Clauses,
diff --git a/clang/lib/AST/StmtPrinter.cpp b/clang/lib/AST/StmtPrinter.cpp
index 6c3294573e9d4..0bacb79ed86e3 100644
--- a/clang/lib/AST/StmtPrinter.cpp
+++ b/clang/lib/AST/StmtPrinter.cpp
@@ -800,6 +800,11 @@ void 
StmtPrinter::VisitOMPInterchangeDirective(OMPInterchangeDirective *Node) {
   PrintOMPExecutableDirective(Node);
 }
 
+void StmtPrinter::VisitOMPFlattenDirective(OMPFlattenDirective *Node) {
+  Indent() << "#pragma omp flatten";
+  PrintOMPExecutableDirective(Node);
+}
+
 void StmtPrinter::VisitOMPSplitDirective(OMPSplitDirective *Node) {
   Indent() << "#pragma omp split";
   PrintOMPExecutableDirective(Node);
diff --git a/clang/lib/AST/StmtProfile.cpp b/clang/lib/AST/StmtProfile.cpp
index 90eab530e0c2e..de7ed322d7424 100644
--- a/clang/lib/AST/StmtProfile.cpp
+++ b/clang/lib/AST/StmtProfile.cpp
@@ -1071,6 +1071,10 @@ void StmtProfiler::VisitOMPInterchangeDirective(
   VisitOMPCanonicalLoopNestTransformationDirective(S);
 }
 
+void StmtProfiler::VisitOMPFlattenDirective(const OMPFlattenDirective *S) {
+  VisitOMPCanonicalLoopNestTransformationDirective(S);
+}
+
 void StmtProfiler::VisitOMPSplitDirective(const OMPSplitDirective *S) {
   VisitOMPCanonicalLoopNestTransformationDirective(S);
 }
diff --git a/clang/lib/Basic/OpenMPKinds.cpp b/clang/lib/Basic/OpenMPKinds.cpp
index 675d86349c933..57e0347f8ea77 100644
--- a/clang/lib/Basic/OpenMPKinds.cpp
+++ b/clang/lib/Basic/OpenMPKinds.cpp
@@ -818,7 +818,7 @@ bool 
clang::isOpenMPCanonicalLoopNestTransformationDirective(
     OpenMPDirectiveKind DKind) {
   return DKind == OMPD_tile || DKind == OMPD_unroll || DKind == OMPD_reverse ||
          DKind == OMPD_split || DKind == OMPD_interchange ||
-         DKind == OMPD_stripe;
+         DKind == OMPD_stripe || DKind == OMPD_flatten;
 }
 
 bool clang::isOpenMPCanonicalLoopSequenceTransformationDirective(

_______________________________________________
cfe-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to