[OpenACC][CIR] Handle firstprivate bounds recipe lowering (#161873)

These work the same as the other two (private and reduction) except that
the expression for the 'init' is a copy instead of a default/value init,
and in a separate region. This patch gets all of that correct, and
ensures we generate these as expected.

There is a little extra work to make sure that the bounds-loop
generation does 2 separate array index operations, otherwise this is
very much like the reduction implementation.

Author: erichkeane

clang/lib/CIR/CodeGen/CIRGenOpenACCRecipe.cpp

@@ -221,10 +221,9 @@ mlir::Value OpenACCRecipeBuilderBase::makeBoundsAlloca(
   return initialAlloca;
 }
 
-mlir::Value
-OpenACCRecipeBuilderBase::createBoundsLoop(mlir::Value subscriptedValue,
-                                           mlir::Value bound,
-                                           mlir::Location loc, bool inverse) {
+std::pair<mlir::Value, mlir::Value> OpenACCRecipeBuilderBase::createBoundsLoop(
+    mlir::Value subscriptedValue, mlir::Value subscriptedValue2,
+    mlir::Value bound, mlir::Location loc, bool inverse) {
   mlir::Operation *bodyInsertLoc;
 
   mlir::Type itrTy = cgf.cgm.convertType(cgf.getContext().UnsignedLongLongTy);
@@ -249,7 +248,6 @@ OpenACCRecipeBuilderBase::createBoundsLoop(mlir::Value subscriptedValue,
 
     return cir::PtrStrideOp::create(builder, loc, eltLoad.getType(), eltLoad,
                                     idxLoad);
-        
   };
 
   auto forStmtBuilder = [&]() {
@@ -303,6 +301,8 @@ OpenACCRecipeBuilderBase::createBoundsLoop(mlir::Value subscriptedValue,
 
           if (subscriptedValue)
             subscriptedValue = doSubscriptOp(subscriptedValue, load);
+          if (subscriptedValue2)
+            subscriptedValue2 = doSubscriptOp(subscriptedValue2, load);
           bodyInsertLoc = builder.createYield(loc);
         },
         /*stepBuilder=*/
@@ -325,7 +325,7 @@ OpenACCRecipeBuilderBase::createBoundsLoop(mlir::Value subscriptedValue,
   // Leave the insertion point to be inside the body, so we can loop over
   // these things.
   builder.setInsertionPoint(bodyInsertLoc);
-  return subscriptedValue;
+  return {subscriptedValue, subscriptedValue2};
 }
 
 mlir::acc::ReductionOperator
@@ -434,7 +434,7 @@ void OpenACCRecipeBuilderBase::createInitRecipe(
     mlir::Location loc, mlir::Location locEnd, SourceRange exprRange,
     mlir::Value mainOp, mlir::Region &recipeInitRegion, size_t numBounds,
     llvm::ArrayRef<QualType> boundTypes, const VarDecl *allocaDecl,
-    QualType origType) {
+    QualType origType, bool emitInitExpr) {
   assert(allocaDecl && "Required recipe variable not set?");
   CIRGenFunction::DeclMapRevertingRAII declMapRAII{cgf, allocaDecl};
 
@@ -464,14 +464,15 @@ void OpenACCRecipeBuilderBase::createInitRecipe(
     // initialize this variable correctly.
     CIRGenFunction::AutoVarEmission tempDeclEmission =
         cgf.emitAutoVarAlloca(*allocaDecl, builder.saveInsertionPoint());
-    cgf.emitAutoVarInit(tempDeclEmission);
+    if (emitInitExpr)
+      cgf.emitAutoVarInit(tempDeclEmission);
   } else {
     mlir::Value alloca = makeBoundsAlloca(
         block, exprRange, loc, allocaDecl->getName(), numBounds, boundTypes);
 
     // If the initializer is trivial, there is nothing to do here, so save
     // ourselves some effort.
-    if (allocaDecl->getInit() &&
+    if (emitInitExpr && allocaDecl->getInit() &&
         (!cgf.isTrivialInitializer(allocaDecl->getInit()) ||
          cgf.getContext().getLangOpts().getTrivialAutoVarInit() !=
              LangOptions::TrivialAutoVarInitKind::Uninitialized))
@@ -484,35 +485,42 @@ void OpenACCRecipeBuilderBase::createInitRecipe(
 
 void OpenACCRecipeBuilderBase::createFirstprivateRecipeCopy(
     mlir::Location loc, mlir::Location locEnd, mlir::Value mainOp,
-    CIRGenFunction::AutoVarEmission tempDeclEmission,
-    mlir::acc::FirstprivateRecipeOp recipe, const VarDecl *varRecipe,
-    const VarDecl *temporary) {
-  mlir::Block *block =
-      createRecipeBlock(recipe.getCopyRegion(), mainOp.getType(), loc,
-                        /*numBounds=*/0, /*isInit=*/false);
-  builder.setInsertionPointToEnd(&recipe.getCopyRegion().back());
+    const VarDecl *allocaDecl, const VarDecl *temporary,
+    mlir::Region &copyRegion, size_t numBounds) {
+  mlir::Block *block = createRecipeBlock(copyRegion, mainOp.getType(), loc,
+                                         numBounds, /*isInit=*/false);
+  builder.setInsertionPointToEnd(&copyRegion.back());
   CIRGenFunction::LexicalScope ls(cgf, loc, block);
 
-  mlir::BlockArgument fromArg = block->getArgument(0);
-  mlir::BlockArgument toArg = block->getArgument(1);
+  mlir::Value fromArg = block->getArgument(0);
+  mlir::Value toArg = block->getArgument(1);
 
-  mlir::Type elementTy =
-      mlir::cast<cir::PointerType>(mainOp.getType()).getPointee();
+  llvm::MutableArrayRef<mlir::BlockArgument> boundsRange =
+      block->getArguments().drop_front(2);
 
-  // Set the address of the emission to be the argument, so that we initialize
-  // that instead of the variable in the other block.
-  tempDeclEmission.setAllocatedAddress(
-      Address{toArg, elementTy, cgf.getContext().getDeclAlign(varRecipe)});
+  for (mlir::BlockArgument boundArg : llvm::reverse(boundsRange))
+    std::tie(fromArg, toArg) =
+        createBoundsLoop(fromArg, toArg, boundArg, loc, /*inverse=*/false);
+
+  // Set up the 'to' address.
+  mlir::Type elementTy =
+      mlir::cast<cir::PointerType>(toArg.getType()).getPointee();
+  CIRGenFunction::AutoVarEmission tempDeclEmission(*allocaDecl);
   tempDeclEmission.emittedAsOffload = true;
+  tempDeclEmission.setAllocatedAddress(
+      Address{toArg, elementTy, cgf.getContext().getDeclAlign(allocaDecl)});
 
+  // Set up the 'from' address from the temporary.
   CIRGenFunction::DeclMapRevertingRAII declMapRAII{cgf, temporary};
   cgf.setAddrOfLocalVar(
       temporary,
-      Address{fromArg, elementTy, cgf.getContext().getDeclAlign(varRecipe)});
-
+      Address{fromArg, elementTy, cgf.getContext().getDeclAlign(allocaDecl)});
   cgf.emitAutoVarInit(tempDeclEmission);
+
+  builder.setInsertionPointToEnd(&copyRegion.back());
   mlir::acc::YieldOp::create(builder, locEnd);
 }
+
 // This function generates the 'combiner' section for a reduction recipe. Note
 // that this function is not 'insertion point' clean, in that it alters the
 // insertion point to be inside of the 'combiner' section of the recipe, but

clang/lib/CIR/CodeGen/CIRGenOpenACCRecipe.h

@@ -49,14 +49,16 @@ class OpenACCRecipeBuilderBase {
   // Creates a loop through an 'acc.bounds', leaving the 'insertion' point to be
   // the inside of the loop body. Traverses LB->UB UNLESS `inverse` is set.
   // Returns the 'subscriptedValue' changed with the new bounds subscript.
+  std::pair<mlir::Value, mlir::Value>
+  createBoundsLoop(mlir::Value subscriptedValue, mlir::Value subscriptedValue2,
+                   mlir::Value bound, mlir::Location loc, bool inverse);
+
   mlir::Value createBoundsLoop(mlir::Value subscriptedValue, mlir::Value bound,
-                               mlir::Location loc, bool inverse);
+                               mlir::Location loc, bool inverse) {
+    return createBoundsLoop(subscriptedValue, {}, bound, loc, inverse).first;
+  }
+
   mlir::acc::ReductionOperator convertReductionOp(OpenACCReductionOperator op);
-  void createFirstprivateRecipeCopy(
-      mlir::Location loc, mlir::Location locEnd, mlir::Value mainOp,
-      CIRGenFunction::AutoVarEmission tempDeclEmission,
-      mlir::acc::FirstprivateRecipeOp recipe, const VarDecl *varRecipe,
-      const VarDecl *temporary);
 
   // This function generates the 'combiner' section for a reduction recipe. Note
   // that this function is not 'insertion point' clean, in that it alters the
@@ -66,11 +68,19 @@ class OpenACCRecipeBuilderBase {
                                      mlir::Value mainOp,
                                      mlir::acc::ReductionRecipeOp recipe,
                                      size_t numBounds);
+
   void createInitRecipe(mlir::Location loc, mlir::Location locEnd,
                         SourceRange exprRange, mlir::Value mainOp,
                         mlir::Region &recipeInitRegion, size_t numBounds,
                         llvm::ArrayRef<QualType> boundTypes,
-                        const VarDecl *allocaDecl, QualType origType);
+                        const VarDecl *allocaDecl, QualType origType,
+                        bool emitInitExpr);
+
+  void createFirstprivateRecipeCopy(mlir::Location loc, mlir::Location locEnd,
+                                    mlir::Value mainOp,
+                                    const VarDecl *allocaDecl,
+                                    const VarDecl *temporary,
+                                    mlir::Region &copyRegion, size_t numBounds);
 
   void createRecipeDestroySection(mlir::Location loc, mlir::Location locEnd,
                                   mlir::Value mainOp, CharUnits alignment,
@@ -150,63 +160,6 @@ class OpenACCRecipeBuilder : OpenACCRecipeBuilderBase {
     return recipeName;
   }
 
-  // Create the 'init' section of the recipe, including the 'copy' section for
-  // 'firstprivate'.  Note that this function is not 'insertion point' clean, in
-  // that it alters the insertion point to be inside of the 'destroy' section of
-  // the recipe, but doesn't restore it aftewards.
-  void createRecipeInitCopy(mlir::Location loc, mlir::Location locEnd,
-                            SourceRange exprRange, mlir::Value mainOp,
-                            RecipeTy recipe, const VarDecl *varRecipe,
-                            const VarDecl *temporary) {
-    // TODO: OpenACC: when we get the 'pointer' variants for
-    // firstprivate/reduction, this probably should be removed/split into
-    // functions for the BuilderBase.
-    assert(varRecipe && "Required recipe variable not set?");
-
-    CIRGenFunction::AutoVarEmission tempDeclEmission{
-        CIRGenFunction::AutoVarEmission::invalid()};
-    CIRGenFunction::DeclMapRevertingRAII declMapRAII{cgf, varRecipe};
-
-    // Do the 'init' section of the recipe IR, which does an alloca, then the
-    // initialization (except for firstprivate).
-    mlir::Block *block =
-        createRecipeBlock(recipe.getInitRegion(), mainOp.getType(), loc,
-                          /*numBounds=*/0, /*isInit=*/true);
-    builder.setInsertionPointToEnd(&recipe.getInitRegion().back());
-    CIRGenFunction::LexicalScope ls(cgf, loc, block);
-
-    tempDeclEmission =
-        cgf.emitAutoVarAlloca(*varRecipe, builder.saveInsertionPoint());
-
-    // 'firstprivate' doesn't do its initialization in the 'init' section,
-    // instead it does it in the 'copy' section.  SO, only do 'init' here for
-    // reduction.
-    if constexpr (std::is_same_v<RecipeTy, mlir::acc::ReductionRecipeOp>) {
-      // Unlike Private, the recipe here is always required as it has to do
-      // init, not just 'default' init.
-      if (!varRecipe->getInit())
-        cgf.cgm.errorNYI(exprRange, "reduction init recipe");
-      cgf.emitAutoVarInit(tempDeclEmission);
-    }
-
-    mlir::acc::YieldOp::create(builder, locEnd);
-
-    if constexpr (std::is_same_v<RecipeTy, mlir::acc::FirstprivateRecipeOp>) {
-      if (!varRecipe->getInit()) {
-        // If we don't have any initialization recipe, we failed during Sema to
-        // initialize this correctly. If we disable the
-        // Sema::TentativeAnalysisScopes in SemaOpenACC::CreateInitRecipe, it'll
-        // emit an error to tell us.  However, emitting those errors during
-        // production is a violation of the standard, so we cannot do them.
-        cgf.cgm.errorNYI(
-            exprRange, "firstprivate copy-init recipe not properly generated");
-      }
-
-      createFirstprivateRecipeCopy(loc, locEnd, mainOp, tempDeclEmission,
-                                   recipe, varRecipe, temporary);
-    }
-  }
-
 public:
   OpenACCRecipeBuilder(CIRGen::CIRGenFunction &cgf,
                        CIRGen::CIRGenBuilderTy &builder)
@@ -221,19 +174,6 @@ class OpenACCRecipeBuilder : OpenACCRecipeBuilderBase {
                BuiltinType::ArraySection) &&
            "array section shouldn't make it to recipe creation");
 
-    // TODO: OpenACC: This is a bit of a hackery to get this to not change for
-    // the non-private recipes. This will be removed soon, when we get this
-    // 'right' for firstprivate and reduction.
-    if constexpr (std::is_same_v<RecipeTy, mlir::acc::FirstprivateRecipeOp>) {
-      if (numBounds) {
-        cgf.cgm.errorNYI(varRef->getSourceRange(),
-                         "firstprivate-init with bounds");
-      }
-      boundTypes = {};
-      numBounds = 0;
-      origType = baseType;
-    }
-
     mlir::ModuleOp mod = builder.getBlock()
                              ->getParent()
                              ->template getParentOfType<mlir::ModuleOp>();
@@ -262,21 +202,20 @@ class OpenACCRecipeBuilder : OpenACCRecipeBuilderBase {
     if constexpr (std::is_same_v<RecipeTy, mlir::acc::PrivateRecipeOp>) {
       createInitRecipe(loc, locEnd, varRef->getSourceRange(), mainOp,
                        recipe.getInitRegion(), numBounds, boundTypes, varRecipe,
-                       origType);
+                       origType, /*emitInitExpr=*/true);
     } else if constexpr (std::is_same_v<RecipeTy,
                                         mlir::acc::ReductionRecipeOp>) {
       createInitRecipe(loc, locEnd, varRef->getSourceRange(), mainOp,
                        recipe.getInitRegion(), numBounds, boundTypes, varRecipe,
-                       origType);
+                       origType, /*emitInitExpr=*/true);
       createReductionRecipeCombiner(loc, locEnd, mainOp, recipe, numBounds);
     } else {
       static_assert(std::is_same_v<RecipeTy, mlir::acc::FirstprivateRecipeOp>);
-      // TODO: OpenACC: we probably want this to call createInitRecipe as well,
-      // but do so in a way that omits the 'initialization', so that we can do
-      // it separately, since it belongs in the 'copy' region. It also might
-      // need a way of getting the tempDeclEmission out of it for that purpose.
-      createRecipeInitCopy(loc, locEnd, varRef->getSourceRange(), mainOp,
-                           recipe, varRecipe, temporary);
+      createInitRecipe(loc, locEnd, varRef->getSourceRange(), mainOp,
+                       recipe.getInitRegion(), numBounds, boundTypes, varRecipe,
+                       origType, /*emitInitExpr=*/false);
+      createFirstprivateRecipeCopy(loc, locEnd, mainOp, varRecipe, temporary,
+                                   recipe.getCopyRegion(), numBounds);
     }
 
     if (origType.isDestructedType())

clang/lib/Sema/SemaOpenACC.cpp

@@ -2724,16 +2724,6 @@ Expr *GenerateReductionInitRecipeExpr(ASTContext &Context,
   return InitExpr;
 }
 
-const Expr *StripOffBounds(const Expr *VarExpr) {
-  while (isa_and_present<ArraySectionExpr, ArraySubscriptExpr>(VarExpr)) {
-    if (const auto *AS = dyn_cast<ArraySectionExpr>(VarExpr))
-      VarExpr = AS->getBase()->IgnoreParenImpCasts();
-    else if (const auto *Sub = dyn_cast<ArraySubscriptExpr>(VarExpr))
-      VarExpr = Sub->getBase()->IgnoreParenImpCasts();
-  }
-  return VarExpr;
-}
-
 VarDecl *CreateAllocaDecl(ASTContext &Ctx, DeclContext *DC,
                           SourceLocation BeginLoc, IdentifierInfo *VarName,
                           QualType VarTy) {
@@ -2794,17 +2784,18 @@ OpenACCPrivateRecipe SemaOpenACC::CreatePrivateInitRecipe(const Expr *VarExpr) {
 
 OpenACCFirstPrivateRecipe
 SemaOpenACC::CreateFirstPrivateInitRecipe(const Expr *VarExpr) {
-  // TODO: OpenACC: This shouldn't be necessary, see PrivateInitRecipe
-  VarExpr = StripOffBounds(VarExpr);
-
+  // We don't strip bounds here, so that we are doing our recipe init at the
+  // 'lowest' possible level.  Codegen is going to have to do its own 'looping'.
   if (!VarExpr || VarExpr->getType()->isDependentType())
     return OpenACCFirstPrivateRecipe::Empty();
 
   QualType VarTy =
       VarExpr->getType().getNonReferenceType().getUnqualifiedType();
 
-  // TODO: OpenACC: for arrays/bounds versions, we're going to have to do a
-  // different initializer, but for now we can go ahead with this.
+  // Array sections are special, and we have to treat them that way.
+  if (const auto *ASE =
+          dyn_cast<ArraySectionExpr>(VarExpr->IgnoreParenImpCasts()))
+    VarTy = ArraySectionExpr::getBaseOriginalType(ASE);
 
   VarDecl *AllocaDecl = CreateAllocaDecl(
       getASTContext(), SemaRef.getCurContext(), VarExpr->getBeginLoc(),