[Clang][OpenMP] Support for dispatch construct (Sema & Codegen) support

clang/docs/OpenMPSupport.rst

@@ -316,6 +316,8 @@ implementation.
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | misc                         | dispatch construct and function variant argument adjustment  | :part:`worked on`        | D99537, D99679                                                        |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
+| misc                         | dispatch construct                                           | :part:`worked on`        |                                                                       |
++------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | misc                         | assumes directives                                           | :part:`worked on`        |                                                                       |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | misc                         | assume directive                                             | :good:`done`             |                                                                       |

clang/docs/ReleaseNotes.rst

@@ -537,6 +537,10 @@ Python Binding Changes
 
 OpenMP Support
 --------------
+- Added support for 'omp assume' directive.
+- Added support for 'omp scope' directive.
+- Added support for allocator-modifier in 'allocate' clause.
+- Added support for 'omp dispatch' directive.
 - Added support 'no_openmp_constructs' assumption clause.
 - Added support for 'self_maps' in map and requirement clause.
 - Added support for 'omp stripe' directive.

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -11888,6 +11888,9 @@ def err_omp_clause_requires_dispatch_construct : Error<
   "'%0' clause requires 'dispatch' context selector">;
 def err_omp_append_args_with_varargs : Error<
   "'append_args' is not allowed with varargs functions">;
+def warn_omp_dispatch_clause_novariants_nocontext : Warning<
+  "'nocontext' clause is ignored, only 'novariants' clause is applied">,
+  InGroup<SourceUsesOpenMP>;
 def err_openmp_vla_in_task_untied : Error<
   "variable length arrays are not supported in OpenMP tasking regions with 'untied' clause">;
 def warn_omp_unterminated_declare_target : Warning<

clang/include/clang/Basic/OpenMPKinds.h

@@ -280,6 +280,12 @@ bool isOpenMPTaskLoopDirective(OpenMPDirectiveKind DKind);
 /// parallel', otherwise - false.
 bool isOpenMPParallelDirective(OpenMPDirectiveKind DKind);
 
+/// Checks if the specified directive is a dispatch-kind directive.
+/// \param DKind Specified directive.
+/// \return true - the directive is a dispatch-like directive like 'omp
+/// dispatch', otherwise - false.
+bool isOpenMPDispatchDirective(OpenMPDirectiveKind DKind);
+
 /// Checks if the specified directive is a target code offload directive.
 /// \param DKind Specified directive.
 /// \return true - the directive is a target code offload directive like

clang/include/clang/Sema/SemaOpenMP.h

@@ -1498,6 +1498,10 @@ class SemaOpenMP : public SemaBase {
                                            : OMPDeclareVariantScopes.back().TI;
   }
 
+  void annotateAStmt(const ASTContext &Context, Stmt *StmtP,
+                     SemaOpenMP *SemaPtr, ArrayRef<OMPClause *> &Clauses,
+                     SourceLocation StartLoc);
+
   /// The current `omp begin/end declare variant` scopes.
   SmallVector<OMPDeclareVariantScope, 4> OMPDeclareVariantScopes;
 

clang/lib/Basic/OpenMPKinds.cpp

@@ -623,6 +623,11 @@ bool clang::isOpenMPParallelDirective(OpenMPDirectiveKind DKind) {
          llvm::is_contained(getLeafConstructs(DKind), OMPD_parallel);
 }
 
+bool clang::isOpenMPDispatchDirective(OpenMPDirectiveKind DKind) {
+  return DKind == OMPD_dispatch ||
+         llvm::is_contained(getLeafConstructs(DKind), OMPD_dispatch);
+}
+
 bool clang::isOpenMPTargetExecutionDirective(OpenMPDirectiveKind DKind) {
   return DKind == OMPD_target ||
          llvm::is_contained(getLeafConstructs(DKind), OMPD_target);
@@ -796,6 +801,10 @@ void clang::getOpenMPCaptureRegions(
       CaptureRegions.push_back(OMPD_task);
       CaptureRegions.push_back(OMPD_target);
       break;
+    case OMPD_dispatch:
+      CaptureRegions.push_back(OMPD_task);
+      CaptureRegions.push_back(OMPD_dispatch);
+      break;
     case OMPD_task:
     case OMPD_target_enter_data:
     case OMPD_target_exit_data:
@@ -819,7 +828,6 @@ void clang::getOpenMPCaptureRegions(
       else
         return true;
       break;
-    case OMPD_dispatch:
     case OMPD_distribute:
     case OMPD_for:
     case OMPD_ordered:

clang/lib/CodeGen/CGStmt.cpp

@@ -424,7 +424,7 @@ void CodeGenFunction::EmitStmt(const Stmt *S, ArrayRef<const Attr *> Attrs) {
     EmitOMPInteropDirective(cast<OMPInteropDirective>(*S));
     break;
   case Stmt::OMPDispatchDirectiveClass:
-    CGM.ErrorUnsupported(S, "OpenMP dispatch directive");
+    EmitOMPDispatchDirective(cast<OMPDispatchDirective>(*S));
     break;
   case Stmt::OMPScopeDirectiveClass:
     EmitOMPScopeDirective(cast<OMPScopeDirective>(*S));

clang/lib/CodeGen/CGStmtOpenMP.cpp

@@ -4528,6 +4528,115 @@ void CodeGenFunction::EmitOMPMasterDirective(const OMPMasterDirective &S) {
   emitMaster(*this, S);
 }
 
+static Expr *getInitialExprFromCapturedExpr(Expr *Cond) {
+
+  Expr *SubExpr = Cond->IgnoreParenImpCasts();
+
+  if (auto *DeclRef = dyn_cast<DeclRefExpr>(SubExpr)) {
+    if (auto *CapturedExprDecl =
+            dyn_cast<OMPCapturedExprDecl>(DeclRef->getDecl())) {
+
+      // Retrieve the initial expression from the captured expression
+      return CapturedExprDecl->getInit();
+    }
+  }
+  return Cond;
+}
+
+static Expr *replaceWithNewTraitsOrDirectCall(Stmt *AssocExpr,
+                                              CallExpr *ReplacementFunction) {
+  Expr *FinalCall = ReplacementFunction;
+
+  if (BinaryOperator *BinaryCopyOpr = dyn_cast<BinaryOperator>(AssocExpr)) {
+    BinaryCopyOpr->setRHS(FinalCall);
+    return BinaryCopyOpr;
+  }
+
+  return FinalCall;
+}
+
+static void transformCallInStmt(Stmt *StmtP) {
+  if (auto *AssocStmt = dyn_cast<CapturedStmt>(StmtP)) {
+    CapturedDecl *CDecl = AssocStmt->getCapturedDecl();
+
+    // Access AnnotateAttr
+    CallExpr *NewCallExpr = nullptr;
+    for (const auto *attr : CDecl->attrs()) {
+      if (const auto *annotateAttr = llvm::dyn_cast<clang::AnnotateAttr>(attr);
+          annotateAttr &&
+          annotateAttr->getAnnotation() == "NoContextInvariant") {
+        NewCallExpr = llvm::dyn_cast<CallExpr>(*annotateAttr->args_begin());
+      }
+    }
+
+    Stmt *CallExprStmt = CDecl->getBody();
+    Stmt *NewCallExprStmt =
+        replaceWithNewTraitsOrDirectCall(CallExprStmt, NewCallExpr);
+    CDecl->setBody(NewCallExprStmt);
+  }
+}
+
+static void EmitIfElse(CodeGenFunction *CGF, Expr *Condition,
+                       Stmt *AssociatedStmt) {
+  llvm::Value *CondValue = CGF->EvaluateExprAsBool(Condition);
+  llvm::BasicBlock *ThenBlock = CGF->createBasicBlock("if.then");
+  llvm::BasicBlock *ElseBlock = CGF->createBasicBlock("if.else");
+  llvm::BasicBlock *MergeBlock = CGF->createBasicBlock("if.end");
+
+  CGF->Builder.CreateCondBr(CondValue, ThenBlock, ElseBlock);
+
+  // Emit the else block.
+  Stmt *ElseStmt = AssociatedStmt;
+  CGF->EmitBlock(ElseBlock);
+  CGF->EmitStmt(ElseStmt);
+  CGF->Builder.CreateBr(MergeBlock);
+
+  // Emit the then block.
+  Stmt *ThenStmt = AssociatedStmt;
+  transformCallInStmt(ThenStmt);
+  CGF->EmitBlock(ThenBlock);
+  CGF->EmitStmt(ThenStmt);
+  CGF->Builder.CreateBr(MergeBlock);
+  CGF->EmitBlock(MergeBlock);
+}
+
+void CodeGenFunction::EmitOMPDispatchDirective(const OMPDispatchDirective &S) {
+  ArrayRef<OMPClause *> Clauses = S.clauses();
+
+  Stmt *AssociatedStmt = const_cast<Stmt *>(S.getAssociatedStmt());
+  if (auto *AssocStmt = dyn_cast<CapturedStmt>(AssociatedStmt))
+    if (auto *InnerCapturedStmt =
+            dyn_cast<CapturedStmt>(AssocStmt->getCapturedStmt())) {
+      AssociatedStmt = InnerCapturedStmt;
+    }
+  CodeGenFunction::CGCapturedStmtInfo CapStmtInfo;
+  if (!CapturedStmtInfo)
+    CapturedStmtInfo = &CapStmtInfo;
+
+  if (!Clauses.empty()) {
+    if (S.hasClausesOfKind<OMPDependClause>())
+      EmitOMPDispatchToTaskwaitDirective(S);
+
+    if (S.hasClausesOfKind<OMPNovariantsClause>() ||
+        S.hasClausesOfKind<OMPNocontextClause>()) {
+      Expr *Condition = nullptr;
+      if (const OMPNovariantsClause *NoVariantsC =
+              OMPExecutableDirective::getSingleClause<OMPNovariantsClause>(
+                  Clauses)) {
+        Condition = getInitialExprFromCapturedExpr(NoVariantsC->getCondition());
+      } else {
+        const OMPNocontextClause *NoContextC =
+            OMPExecutableDirective::getSingleClause<OMPNocontextClause>(
+                Clauses);
+        Condition = getInitialExprFromCapturedExpr(NoContextC->getCondition());
+      }
+      /* OMPC_novariants or OMPC_nocontext present */
+      EmitIfElse(this, Condition, AssociatedStmt);
+    }
+  } else
+    EmitStmt(AssociatedStmt);
+}
+
 static void emitMasked(CodeGenFunction &CGF, const OMPExecutableDirective &S) {
   auto &&CodeGen = [&S](CodeGenFunction &CGF, PrePostActionTy &Action) {
     Action.Enter(CGF);
@@ -5548,6 +5657,15 @@ void CodeGenFunction::EmitOMPBarrierDirective(const OMPBarrierDirective &S) {
   CGM.getOpenMPRuntime().emitBarrierCall(*this, S.getBeginLoc(), OMPD_barrier);
 }
 
+void CodeGenFunction::EmitOMPDispatchToTaskwaitDirective(
+    const OMPDispatchDirective &S) {
+  OMPTaskDataTy Data;
+  // Build list of dependences
+  buildDependences(S, Data);
+  Data.HasNowaitClause = S.hasClausesOfKind<OMPNowaitClause>();
+  CGM.getOpenMPRuntime().emitTaskwaitCall(*this, S.getBeginLoc(), Data);
+}
+
 void CodeGenFunction::EmitOMPTaskwaitDirective(const OMPTaskwaitDirective &S) {
   OMPTaskDataTy Data;
   // Build list of dependences

clang/lib/CodeGen/CodeGenFunction.h

@@ -3676,6 +3676,8 @@ class CodeGenFunction : public CodeGenTypeCache {
   void EmitCXXForRangeStmt(const CXXForRangeStmt &S,
                            ArrayRef<const Attr *> Attrs = {});
 
+  void EmitOMPDispatchToTaskwaitDirective(const OMPDispatchDirective &S);
+
   /// Controls insertion of cancellation exit blocks in worksharing constructs.
   class OMPCancelStackRAII {
     CodeGenFunction &CGF;
@@ -3868,6 +3870,7 @@ class CodeGenFunction : public CodeGenTypeCache {
   void EmitOMPSectionDirective(const OMPSectionDirective &S);
   void EmitOMPSingleDirective(const OMPSingleDirective &S);
   void EmitOMPMasterDirective(const OMPMasterDirective &S);
+  void EmitOMPDispatchDirective(const OMPDispatchDirective &S);
   void EmitOMPMaskedDirective(const OMPMaskedDirective &S);
   void EmitOMPCriticalDirective(const OMPCriticalDirective &S);
   void EmitOMPParallelForDirective(const OMPParallelForDirective &S);