[clang] [llvm] Support for dispatch construct (Sema & Codegen) support. (PR #117904)

[Alexey Bataev via cfe-commits]

================
@@ -5965,6 +5967,263 @@ static bool teamsLoopCanBeParallelFor(Stmt *AStmt, Sema &SemaRef) {
   return Checker.teamsLoopCanBeParallelFor();
 }
 
+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 nullptr;
+}
+
+// Next two functions, cloneAssociatedStmt() &
+// replaceWithNewTraitsOrDirectCall(), are for transforming the call traits.
+// e.g.
+// #pragma omp declare variant(foo_variant_dispatch) match(construct={dispatch})
+// #pragma omp declare variant(foo_variant_allCond) match(user={condition(1)})
+// ..
+//     #pragma omp dispatch nocontext(cond_true)
+//         foo(i, j);
+// is changed to:
+// if (cond_true) {
+//    foo(i,j) // with traits: CodeGen call to foo_variant_allCond(i,j)
+// } else {
+//   #pragma omp dispatch
+//   foo(i,j)  // with traits: CodeGen call to foo_variant_dispatch(i,j)
+// }
+//
+// The next 2 functions, are for:
+// if (cond_true) {
+//   foo(i,j) // with traits: runtime call to foo_variant_allCond(i,j)
+// }
+//
+static Expr *replaceWithNewTraitsOrDirectCall(const ASTContext &Context, Expr *,
+                                              SemaOpenMP *, bool);
+
+static StmtResult cloneAssociatedStmt(const ASTContext &Context, Stmt *StmtP,
+                                      SemaOpenMP *SemaPtr, bool NoContext) {
+  if (CapturedStmt *AssocStmt = dyn_cast<CapturedStmt>(StmtP)) {
+    CapturedDecl *CDecl = AssocStmt->getCapturedDecl();
+    Stmt *AssocExprStmt = AssocStmt->getCapturedStmt();
+    Expr *AssocExpr = dyn_cast<Expr>(AssocExprStmt);
+    Expr *NewCallOrPseudoObjOrBinExpr = replaceWithNewTraitsOrDirectCall(
+        Context, AssocExpr, SemaPtr, NoContext);
+
+    // Copy Current Captured Decl to a New Captured Decl for noting the
+    // Annotation
+    CapturedDecl *NewDecl =
+        CapturedDecl::Create(const_cast<ASTContext &>(Context),
+                             CDecl->getDeclContext(), CDecl->getNumParams());
+    NewDecl->setBody(static_cast<Stmt *>(NewCallOrPseudoObjOrBinExpr));
+    for (unsigned i = 0; i < CDecl->getNumParams(); ++i) {
+      if (i != CDecl->getContextParamPosition())
+        NewDecl->setParam(i, CDecl->getParam(i));
+      else
+        NewDecl->setContextParam(i, CDecl->getContextParam());
+    }
+
+    // Create a New Captured Stmt containing the New Captured Decl
+    SmallVector<CapturedStmt::Capture, 4> Captures;
+    SmallVector<Expr *, 4> CaptureInits;
+    for (auto capture : AssocStmt->captures())
+      Captures.push_back(capture);
+    for (auto capture_init : AssocStmt->capture_inits())
----------------
alexey-bataev wrote:

`for (Expr *CaptureInit : AssocStmt->capture_inits())`

https://github.com/llvm/llvm-project/pull/117904