================
@@ -15871,6 +15881,172 @@ StmtResult 
SemaOpenMP::ActOnOpenMPReverseDirective(Stmt *AStmt,
                                      buildPreInits(Context, PreInits));
 }
 
+/// Build the AST for \#pragma omp split counts(c1, c2, ...).
+///
+/// Splits the single associated loop into N consecutive loops, where N is the
+/// number of count expressions.
+StmtResult SemaOpenMP::ActOnOpenMPSplitDirective(ArrayRef<OMPClause *> Clauses,
+                                                 Stmt *AStmt,
+                                                 SourceLocation StartLoc,
+                                                 SourceLocation EndLoc) {
+  ASTContext &Context = getASTContext();
+  Scope *CurScope = SemaRef.getCurScope();
+
+  // Empty statement should only be possible if there already was an error.
+  if (!AStmt)
+    return StmtError();
+
+  const OMPCountsClause *CountsClause =
+      OMPExecutableDirective::getSingleClause<OMPCountsClause>(Clauses);
+  if (!CountsClause)
+    return StmtError();
+
+  // Split applies to a single loop; check it is transformable and get helpers.
+  constexpr unsigned NumLoops = 1;
+  Stmt *Body = nullptr;
+  SmallVector<OMPLoopBasedDirective::HelperExprs, NumLoops> LoopHelpers(
+      NumLoops);
+  SmallVector<SmallVector<Stmt *>, NumLoops + 1> OriginalInits;
+  if (!checkTransformableLoopNest(OMPD_split, AStmt, NumLoops, LoopHelpers,
+                                  Body, OriginalInits))
+    return StmtError();
+
+  // Delay applying the transformation to when template is completely
+  // instantiated.
+  if (SemaRef.CurContext->isDependentContext())
+    return OMPSplitDirective::Create(Context, StartLoc, EndLoc, Clauses,
+                                     NumLoops, AStmt, nullptr, nullptr);
+
+  assert(LoopHelpers.size() == NumLoops &&
+         "Expecting a single-dimensional loop iteration space");
+  assert(OriginalInits.size() == NumLoops &&
+         "Expecting a single-dimensional loop iteration space");
+  OMPLoopBasedDirective::HelperExprs &LoopHelper = LoopHelpers.front();
+
+  // Find the loop statement.
+  Stmt *LoopStmt = nullptr;
+  collectLoopStmts(AStmt, {LoopStmt});
+
+  // Determine the PreInit declarations.
+  SmallVector<Stmt *> PreInits;
+  addLoopPreInits(Context, LoopHelper, LoopStmt, OriginalInits[0], PreInits);
+
+  // Type and name of the original loop variable; we create one IV per segment
+  // and assign it to the original var so the body sees the same name.
+  auto *IterationVarRef = cast<DeclRefExpr>(LoopHelper.IterationVarRef);
+  QualType IVTy = IterationVarRef->getType();
+  uint64_t IVWidth = Context.getTypeSize(IVTy);
+  auto *OrigVar = cast<DeclRefExpr>(LoopHelper.Counters.front());
+
+  // Iteration variable SourceLocations.
+  SourceLocation OrigVarLoc = OrigVar->getExprLoc();
+  SourceLocation OrigVarLocBegin = OrigVar->getBeginLoc();
+  SourceLocation OrigVarLocEnd = OrigVar->getEndLoc();
+  // Internal variable names.
+  std::string OrigVarName = OrigVar->getNameInfo().getAsString();
+
+  // Collect constant count values from the counts clause
+  SmallVector<uint64_t, 4> CountValues;
+  for (Expr *CountExpr : CountsClause->getCountsRefs()) {
+    if (!CountExpr) {
+      return OMPSplitDirective::Create(Context, StartLoc, EndLoc, Clauses,
+                                       NumLoops, AStmt, nullptr, nullptr);
+    }
+    std::optional<llvm::APSInt> OptVal =
+        CountExpr->getIntegerConstantExpr(Context);
+    if (!OptVal || OptVal->isNegative()) {
+      return OMPSplitDirective::Create(Context, StartLoc, EndLoc, Clauses,
+                                       NumLoops, AStmt, nullptr, nullptr);
+    }
+    CountValues.push_back(OptVal->getZExtValue());
+  }
+
+  if (CountValues.empty()) {
+    Diag(CountsClause->getBeginLoc(), diag::err_omp_unexpected_clause_value)
+        << "at least one non-negative integer expression" << "counts";
+    return StmtError();
+  }
+
+  // Cumulative segment starts: Starts[0]=0,
+  // Starts[j]=Starts[j-1]+CountValues[j-1]. Example: CountValues [3,5,2] →
+  // Starts [0,3,8,10]. Segment k runs [Starts[k], Starts[k+1]).
+  SmallVector<uint64_t, 4> Starts;
+  Starts.push_back(0);
+  for (size_t j = 0; j < CountValues.size(); ++j)
+    Starts.push_back(Starts.back() + CountValues[j]);
+
+  size_t NumSegments = CountValues.size();
+  SmallVector<Stmt *, 4> SplitLoops;
+
+  for (size_t Seg = 0; Seg < NumSegments; ++Seg) {
+    uint64_t StartVal = Starts[Seg];
+    uint64_t EndVal = Starts[Seg + 1];
+
+    // Segment IV: .split.iv.<Seg>.<OrigVarName>, init to StartVal, bound by
+    // EndVal.
+    SmallString<64> IVName(".split.iv.");
+    IVName += Twine(Seg).str();
+    IVName += ".";
+    IVName += OrigVarName;
----------------
Meinersbur wrote:

```suggestion
    IVName += (Twine(Seg) + "." + OrigVarName).str();
```
This is what a Twine is for

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

Reply via email to