[OpenMP] Sema and parsing for 'target teams distribute' pragma

This patch is to implement sema and parsing for 'target teams distribute' pragma.

Differential Revision: https://reviews.llvm.org/D28015

llvm-svn: 290508

clang/include/clang-c/Index.h

@@ -2358,7 +2358,11 @@ enum CXCursorKind {
    */
   CXCursor_OMPTargetTeamsDirective = 275,
 
-  CXCursor_LastStmt = CXCursor_OMPTargetTeamsDirective,
+  /** \brief OpenMP target teams distribute directive.
+   */
+  CXCursor_OMPTargetTeamsDistributeDirective = 276,
+
+  CXCursor_LastStmt = CXCursor_OMPTargetTeamsDistributeDirective,
 
   /**
    * \brief Cursor that represents the translation unit itself.

clang/include/clang/AST/RecursiveASTVisitor.h

@@ -2663,6 +2663,9 @@ DEF_TRAVERSE_STMT(OMPTeamsDistributeParallelForDirective,
 DEF_TRAVERSE_STMT(OMPTargetTeamsDirective,
                   { TRY_TO(TraverseOMPExecutableDirective(S)); })
 
+DEF_TRAVERSE_STMT(OMPTargetTeamsDistributeDirective,
+                  { TRY_TO(TraverseOMPExecutableDirective(S)); })
+
 // OpenMP clauses.
 template <typename Derived>
 bool RecursiveASTVisitor<Derived>::TraverseOMPClause(OMPClause *C) {

clang/include/clang/AST/StmtOpenMP.h

@@ -3498,6 +3498,73 @@ class OMPTargetTeamsDirective final : public OMPExecutableDirective {
   }
 };
 
+/// This represents '#pragma omp target teams distribute' combined directive.
+///
+/// \code
+/// #pragma omp target teams distribute private(x)
+/// \endcode
+/// In this example directive '#pragma omp target teams distribute' has clause
+/// 'private' with the variables 'x'
+///
+class OMPTargetTeamsDistributeDirective final : public OMPLoopDirective {
+  friend class ASTStmtReader;
+
+  /// Build directive with the given start and end location.
+  ///
+  /// \param StartLoc Starting location of the directive kind.
+  /// \param EndLoc Ending location of the directive.
+  /// \param CollapsedNum Number of collapsed nested loops.
+  /// \param NumClauses Number of clauses.
+  ///
+  OMPTargetTeamsDistributeDirective(SourceLocation StartLoc,
+                                    SourceLocation EndLoc,
+                                    unsigned CollapsedNum, unsigned NumClauses)
+      : OMPLoopDirective(this, OMPTargetTeamsDistributeDirectiveClass,
+                         OMPD_target_teams_distribute, StartLoc, EndLoc,
+                         CollapsedNum, NumClauses) {}
+
+  /// Build an empty directive.
+  ///
+  /// \param CollapsedNum Number of collapsed nested loops.
+  /// \param NumClauses Number of clauses.
+  ///
+  explicit OMPTargetTeamsDistributeDirective(unsigned CollapsedNum,
+                                             unsigned NumClauses)
+      : OMPLoopDirective(this, OMPTargetTeamsDistributeDirectiveClass,
+                         OMPD_target_teams_distribute, SourceLocation(),
+                         SourceLocation(), CollapsedNum, NumClauses) {}
+
+public:
+  /// Creates directive with a list of \a Clauses.
+  ///
+  /// \param C AST context.
+  /// \param StartLoc Starting location of the directive kind.
+  /// \param EndLoc Ending Location of the directive.
+  /// \param CollapsedNum Number of collapsed loops.
+  /// \param Clauses List of clauses.
+  /// \param AssociatedStmt Statement, associated with the directive.
+  /// \param Exprs Helper expressions for CodeGen.
+  ///
+  static OMPTargetTeamsDistributeDirective *
+  Create(const ASTContext &C, SourceLocation StartLoc, SourceLocation EndLoc,
+         unsigned CollapsedNum, ArrayRef<OMPClause *> Clauses,
+         Stmt *AssociatedStmt, const HelperExprs &Exprs);
+
+  /// Creates an empty directive with the place for \a NumClauses clauses.
+  ///
+  /// \param C AST context.
+  /// \param CollapsedNum Number of collapsed nested loops.
+  /// \param NumClauses Number of clauses.
+  ///
+  static OMPTargetTeamsDistributeDirective *
+  CreateEmpty(const ASTContext &C, unsigned NumClauses, unsigned CollapsedNum,
+              EmptyShell);
+
+  static bool classof(const Stmt *T) {
+    return T->getStmtClass() == OMPTargetTeamsDistributeDirectiveClass;
+  }
+};
+
 } // end namespace clang
 
 #endif

clang/include/clang/Basic/OpenMPKinds.def

@@ -156,6 +156,9 @@
 #ifndef OPENMP_TARGET_TEAMS_CLAUSE
 #define OPENMP_TARGET_TEAMS_CLAUSE(Name)
 #endif
+#ifndef OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE
+#define OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(Name)
+#endif
 
 // OpenMP directives.
 OPENMP_DIRECTIVE(threadprivate)
@@ -206,6 +209,7 @@ OPENMP_DIRECTIVE_EXT(teams_distribute_simd, "teams distribute simd")
 OPENMP_DIRECTIVE_EXT(teams_distribute_parallel_for_simd, "teams distribute parallel for simd")
 OPENMP_DIRECTIVE_EXT(teams_distribute_parallel_for, "teams distribute parallel for")
 OPENMP_DIRECTIVE_EXT(target_teams, "target teams")
+OPENMP_DIRECTIVE_EXT(target_teams_distribute, "target teams distribute")
 
 // OpenMP clauses.
 OPENMP_CLAUSE(if, OMPIfClause)
@@ -743,6 +747,25 @@ OPENMP_TARGET_TEAMS_CLAUSE(reduction)
 OPENMP_TARGET_TEAMS_CLAUSE(num_teams)
 OPENMP_TARGET_TEAMS_CLAUSE(thread_limit)
 
+// Clauses allowed for OpenMP directive 'target teams distribute'.
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(if)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(device)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(map)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(private)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(nowait)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(depend)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(defaultmap)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(firstprivate)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(is_device_ptr)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(default)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(shared)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(reduction)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(num_teams)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(thread_limit)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(lastprivate)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(collapse)
+OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(dist_schedule)
+
 #undef OPENMP_TASKLOOP_SIMD_CLAUSE
 #undef OPENMP_TASKLOOP_CLAUSE
 #undef OPENMP_LINEAR_KIND
@@ -791,3 +814,4 @@ OPENMP_TARGET_TEAMS_CLAUSE(thread_limit)
 #undef OPENMP_TEAMS_DISTRIBUTE_PARALLEL_FOR_SIMD_CLAUSE
 #undef OPENMP_TEAMS_DISTRIBUTE_PARALLEL_FOR_CLAUSE
 #undef OPENMP_TARGET_TEAMS_CLAUSE
+#undef OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE

clang/include/clang/Basic/StmtNodes.td

@@ -241,3 +241,4 @@ def OMPTeamsDistributeSimdDirective : DStmt<OMPLoopDirective>;
 def OMPTeamsDistributeParallelForSimdDirective : DStmt<OMPLoopDirective>;
 def OMPTeamsDistributeParallelForDirective : DStmt<OMPLoopDirective>;
 def OMPTargetTeamsDirective : DStmt<OMPExecutableDirective>;
+def OMPTargetTeamsDistributeDirective : DStmt<OMPLoopDirective>;

clang/include/clang/Sema/Sema.h

@@ -8471,6 +8471,12 @@ class Sema {
                                              Stmt *AStmt,
                                              SourceLocation StartLoc,
                                              SourceLocation EndLoc);
+  /// Called on well-formed '\#pragma omp target teams distribute' after parsing
+  /// of the associated statement.
+  StmtResult ActOnOpenMPTargetTeamsDistributeDirective(
+      ArrayRef<OMPClause *> Clauses, Stmt *AStmt, SourceLocation StartLoc,
+      SourceLocation EndLoc,
+      llvm::DenseMap<ValueDecl *, Expr *> &VarsWithImplicitDSA);
 
   /// Checks correctness of linear modifiers.
   bool CheckOpenMPLinearModifier(OpenMPLinearClauseKind LinKind,

clang/include/clang/Serialization/ASTBitCodes.h

@@ -1514,6 +1514,7 @@ namespace clang {
       STMT_OMP_TEAMS_DISTRIBUTE_PARALLEL_FOR_SIMD_DIRECTIVE,
       STMT_OMP_TEAMS_DISTRIBUTE_PARALLEL_FOR_DIRECTIVE,
       STMT_OMP_TARGET_TEAMS_DIRECTIVE,
+      STMT_OMP_TARGET_TEAMS_DISTRIBUTE_DIRECTIVE,
       EXPR_OMP_ARRAY_SECTION,
 
       // ARC

clang/lib/AST/StmtOpenMP.cpp

@@ -1543,3 +1543,58 @@ OMPTargetTeamsDirective::CreateEmpty(const ASTContext &C, unsigned NumClauses,
       C.Allocate(Size + sizeof(OMPClause *) * NumClauses + sizeof(Stmt *));
   return new (Mem) OMPTargetTeamsDirective(NumClauses);
 }
+
+OMPTargetTeamsDistributeDirective *OMPTargetTeamsDistributeDirective::Create(
+    const ASTContext &C, SourceLocation StartLoc, SourceLocation EndLoc,
+    unsigned CollapsedNum, ArrayRef<OMPClause *> Clauses, Stmt *AssociatedStmt,
+    const HelperExprs &Exprs) {
+  auto Size = llvm::alignTo(sizeof(OMPTargetTeamsDistributeDirective),
+                            alignof(OMPClause *));
+  void *Mem = C.Allocate(
+      Size + sizeof(OMPClause *) * Clauses.size() +
+      sizeof(Stmt *) *
+          numLoopChildren(CollapsedNum, OMPD_target_teams_distribute));
+  OMPTargetTeamsDistributeDirective *Dir =
+      new (Mem) OMPTargetTeamsDistributeDirective(StartLoc, EndLoc, CollapsedNum,
+                                                  Clauses.size());
+  Dir->setClauses(Clauses);
+  Dir->setAssociatedStmt(AssociatedStmt);
+  Dir->setIterationVariable(Exprs.IterationVarRef);
+  Dir->setLastIteration(Exprs.LastIteration);
+  Dir->setCalcLastIteration(Exprs.CalcLastIteration);
+  Dir->setPreCond(Exprs.PreCond);
+  Dir->setCond(Exprs.Cond);
+  Dir->setInit(Exprs.Init);
+  Dir->setInc(Exprs.Inc);
+  Dir->setIsLastIterVariable(Exprs.IL);
+  Dir->setLowerBoundVariable(Exprs.LB);
+  Dir->setUpperBoundVariable(Exprs.UB);
+  Dir->setStrideVariable(Exprs.ST);
+  Dir->setEnsureUpperBound(Exprs.EUB);
+  Dir->setNextLowerBound(Exprs.NLB);
+  Dir->setNextUpperBound(Exprs.NUB);
+  Dir->setNumIterations(Exprs.NumIterations);
+  Dir->setPrevLowerBoundVariable(Exprs.PrevLB);
+  Dir->setPrevUpperBoundVariable(Exprs.PrevUB);
+  Dir->setCounters(Exprs.Counters);
+  Dir->setPrivateCounters(Exprs.PrivateCounters);
+  Dir->setInits(Exprs.Inits);
+  Dir->setUpdates(Exprs.Updates);
+  Dir->setFinals(Exprs.Finals);
+  Dir->setPreInits(Exprs.PreInits);
+  return Dir;
+}
+
+OMPTargetTeamsDistributeDirective *
+OMPTargetTeamsDistributeDirective::CreateEmpty(const ASTContext &C,
+                                               unsigned NumClauses,
+                                               unsigned CollapsedNum,
+                                               EmptyShell) {
+  auto Size = llvm::alignTo(sizeof(OMPTargetTeamsDistributeDirective),
+                            alignof(OMPClause *));
+  void *Mem = C.Allocate(
+      Size + sizeof(OMPClause *) * NumClauses +
+      sizeof(Stmt *) *
+           numLoopChildren(CollapsedNum, OMPD_target_teams_distribute));
+  return new (Mem) OMPTargetTeamsDistributeDirective(CollapsedNum, NumClauses);
+}

clang/lib/AST/StmtPrinter.cpp

@@ -1232,6 +1232,12 @@ void StmtPrinter::VisitOMPTargetTeamsDirective(OMPTargetTeamsDirective *Node) {
   PrintOMPExecutableDirective(Node);
 }
 
+void StmtPrinter::VisitOMPTargetTeamsDistributeDirective(
+    OMPTargetTeamsDistributeDirective *Node) {
+  Indent() << "#pragma omp target teams distribute ";
+  PrintOMPExecutableDirective(Node);
+}
+
 //===----------------------------------------------------------------------===//
 //  Expr printing methods.
 //===----------------------------------------------------------------------===//

clang/lib/AST/StmtProfile.cpp

@@ -753,6 +753,11 @@ void StmtProfiler::VisitOMPTargetTeamsDirective(
   VisitOMPExecutableDirective(S);
 }
 
+void StmtProfiler::VisitOMPTargetTeamsDistributeDirective(
+    const OMPTargetTeamsDistributeDirective *S) {
+  VisitOMPLoopDirective(S);
+}
+
 void StmtProfiler::VisitExpr(const Expr *S) {
   VisitStmt(S);
 }

clang/lib/Basic/OpenMPKinds.cpp

@@ -665,6 +665,16 @@ bool clang::isAllowedClauseForDirective(OpenMPDirectiveKind DKind,
 #define OPENMP_TARGET_TEAMS_CLAUSE(Name)                                       \
   case OMPC_##Name:                                                            \
     return true;
+#include "clang/Basic/OpenMPKinds.def"
+    default:
+      break;
+    }
+    break;
+  case OMPD_target_teams_distribute:
+    switch (CKind) {
+#define OPENMP_TARGET_TEAMS_DISTRIBUTE_CLAUSE(Name)                            \
+  case OMPC_##Name:                                                            \
+    return true;
 #include "clang/Basic/OpenMPKinds.def"
     default:
       break;
@@ -699,8 +709,8 @@ bool clang::isOpenMPLoopDirective(OpenMPDirectiveKind DKind) {
          DKind == OMPD_teams_distribute ||
          DKind == OMPD_teams_distribute_simd ||
          DKind == OMPD_teams_distribute_parallel_for_simd ||
-         DKind == OMPD_teams_distribute_parallel_for;
-  // TODO add next directives.
+         DKind == OMPD_teams_distribute_parallel_for ||
+         DKind == OMPD_target_teams_distribute;
 }
 
 bool clang::isOpenMPWorksharingDirective(OpenMPDirectiveKind DKind) {
@@ -737,7 +747,7 @@ bool clang::isOpenMPTargetExecutionDirective(OpenMPDirectiveKind DKind) {
   return DKind == OMPD_target || DKind == OMPD_target_parallel ||
          DKind == OMPD_target_parallel_for || 
          DKind == OMPD_target_parallel_for_simd || DKind == OMPD_target_simd ||
-         DKind == OMPD_target_teams;
+         DKind == OMPD_target_teams || DKind == OMPD_target_teams_distribute;
 }
 
 bool clang::isOpenMPTargetDataManagementDirective(OpenMPDirectiveKind DKind) {
@@ -754,7 +764,7 @@ bool clang::isOpenMPNestingTeamsDirective(OpenMPDirectiveKind DKind) {
 
 bool clang::isOpenMPTeamsDirective(OpenMPDirectiveKind DKind) {
   return isOpenMPNestingTeamsDirective(DKind) ||
-         DKind == OMPD_target_teams;
+         DKind == OMPD_target_teams || DKind == OMPD_target_teams_distribute;
 }
 
 bool clang::isOpenMPSimdDirective(OpenMPDirectiveKind DKind) {
@@ -778,8 +788,8 @@ bool clang::isOpenMPDistributeDirective(OpenMPDirectiveKind Kind) {
   return isOpenMPNestingDistributeDirective(Kind) ||
          Kind == OMPD_teams_distribute || Kind == OMPD_teams_distribute_simd ||
          Kind == OMPD_teams_distribute_parallel_for_simd ||
-         Kind == OMPD_teams_distribute_parallel_for;
-  // TODO add next directives.
+         Kind == OMPD_teams_distribute_parallel_for ||
+         Kind == OMPD_target_teams_distribute;
 }
 
 bool clang::isOpenMPPrivate(OpenMPClauseKind Kind) {
@@ -802,5 +812,6 @@ bool clang::isOpenMPLoopBoundSharingDirective(OpenMPDirectiveKind Kind) {
          Kind == OMPD_distribute_simd || Kind == OMPD_teams_distribute ||
          Kind == OMPD_teams_distribute_simd ||
          Kind == OMPD_teams_distribute_parallel_for_simd ||
-         Kind == OMPD_teams_distribute_parallel_for;
+         Kind == OMPD_teams_distribute_parallel_for ||
+         Kind == OMPD_target_teams_distribute;
 }

clang/lib/CodeGen/CGStmt.cpp

@@ -318,6 +318,10 @@ void CodeGenFunction::EmitStmt(const Stmt *S) {
   case Stmt::OMPTargetTeamsDirectiveClass:
     EmitOMPTargetTeamsDirective(cast<OMPTargetTeamsDirective>(*S));
     break;
+  case Stmt::OMPTargetTeamsDistributeDirectiveClass:
+    EmitOMPTargetTeamsDistributeDirective(
+        cast<OMPTargetTeamsDistributeDirective>(*S));
+    break;
   }
 }
 

clang/lib/CodeGen/CGStmtOpenMP.cpp

@@ -2011,6 +2011,14 @@ void CodeGenFunction::EmitOMPTargetTeamsDirective(
       });
 }
 
+void CodeGenFunction::EmitOMPTargetTeamsDistributeDirective(
+    const OMPTargetTeamsDistributeDirective &S) {
+  CGM.getOpenMPRuntime().emitInlinedDirective(*this, OMPD_target_teams_distribute,
+      [&S](CodeGenFunction &CGF, PrePostActionTy &) {
+        CGF.EmitStmt(cast<CapturedStmt>(S.getAssociatedStmt())->getCapturedStmt());
+      });
+}
+
 /// \brief Emit a helper variable and return corresponding lvalue.
 static LValue EmitOMPHelperVar(CodeGenFunction &CGF,
                                const DeclRefExpr *Helper) {

clang/lib/CodeGen/CodeGenFunction.h

@@ -2695,6 +2695,8 @@ class CodeGenFunction : public CodeGenTypeCache {
   void EmitOMPTeamsDistributeParallelForDirective(
       const OMPTeamsDistributeParallelForDirective &S);
   void EmitOMPTargetTeamsDirective(const OMPTargetTeamsDirective &S);
+  void EmitOMPTargetTeamsDistributeDirective(
+      const OMPTargetTeamsDistributeDirective &S);
 
   /// Emit outlined function for the target directive.
   static std::pair<llvm::Function * /*OutlinedFn*/,

clang/lib/Parse/ParseOpenMP.cpp

@@ -115,7 +115,8 @@ static OpenMPDirectiveKind ParseOpenMPDirectiveKind(Parser &P) {
     { OMPD_teams_distribute, OMPD_parallel, OMPD_teams_distribute_parallel },
     { OMPD_teams_distribute_parallel, OMPD_for, OMPD_teams_distribute_parallel_for },
     { OMPD_teams_distribute_parallel_for, OMPD_simd, OMPD_teams_distribute_parallel_for_simd },
-    { OMPD_target, OMPD_teams, OMPD_target_teams }
+    { OMPD_target, OMPD_teams, OMPD_target_teams },
+    { OMPD_target_teams, OMPD_distribute, OMPD_target_teams_distribute }
   };
   enum { CancellationPoint = 0, DeclareReduction = 1, TargetData = 2 };
   auto Tok = P.getCurToken();
@@ -752,6 +753,7 @@ Parser::DeclGroupPtrTy Parser::ParseOpenMPDeclarativeDirectiveWithExtDecl(
   case OMPD_teams_distribute_parallel_for_simd:
   case OMPD_teams_distribute_parallel_for:
   case OMPD_target_teams:
+  case OMPD_target_teams_distribute:
     Diag(Tok, diag::err_omp_unexpected_directive)
         << getOpenMPDirectiveName(DKind);
     break;
@@ -788,7 +790,8 @@ Parser::DeclGroupPtrTy Parser::ParseOpenMPDeclarativeDirectiveWithExtDecl(
 ///         'target parallel for simd' | 'target simd' |
 ///         'teams distribute' | 'teams distribute simd' |
 ///         'teams distribute parallel for simd' |
-///         'teams distribute parallel for' | 'target teams' {clause}
+///         'teams distribute parallel for' | 'target teams'
+///         'target teams distribute' {clause}
 ///         annot_pragma_openmp_end
 ///
 StmtResult Parser::ParseOpenMPDeclarativeOrExecutableDirective(
@@ -902,7 +905,8 @@ StmtResult Parser::ParseOpenMPDeclarativeOrExecutableDirective(
   case OMPD_teams_distribute_simd:
   case OMPD_teams_distribute_parallel_for_simd:
   case OMPD_teams_distribute_parallel_for:
-  case OMPD_target_teams: {
+  case OMPD_target_teams:
+  case OMPD_target_teams_distribute: {
     ConsumeToken();
     // Parse directive name of the 'critical' directive if any.
     if (DKind == OMPD_critical) {