[flang][openacc] Populate init and copy region of firstprivate recipe for simple type

[clementval]

Summary

Similiar to D154259 for private recipe, this patch makes use of the same code
for the init region of the firstprivate recipe and populate the copy region for
trivial types and arrays.

[razvanlupusoru]

Looks great! Thank you!

flang/lib/Lower/OpenACC.cpp

if constexpr (std::is_same_v<ExitOp, mlir::acc::CopyoutOp> ||
       varPtr = entryOp.getVarPtr();
     builder.create<ExitOp>(entryOp.getLoc(), entryOp.getAccPtr(), varPtr,
                            entryOp.getBounds(), entryOp.getDataClause(),
                            structured, implicit,
                            builder.getStringAttr(*entryOp.getName()));
   }
 }
 
+template <typename RecipeOp>
+static void genPrivateLikeInitRegion(mlir::OpBuilder &builder, RecipeOp recipe,
+                                     mlir::Type ty, mlir::Location loc) {
+  mlir::Value retVal = recipe.getInitRegion().front().getArgument(0);
+  if (auto refTy = mlir::dyn_cast_or_null<fir::ReferenceType>(ty)) {
+    if (fir::isa_trivial(refTy.getEleTy()))
+      retVal = builder.create<fir::AllocaOp>(loc, refTy.getEleTy());
+    else if (auto seqTy =
+                 mlir::dyn_cast_or_null<fir::SequenceType>(refTy.getEleTy())) {
+      if (seqTy.hasDynamicExtents())
+        TODO(loc, "private recipe of array with dynamic extents");
+      if (fir::isa_trivial(seqTy.getEleTy()))
+        retVal = builder.create<fir::AllocaOp>(loc, seqTy);
+    }
+  }
+  builder.create<mlir::acc::YieldOp>(loc, retVal);
+}
+
 mlir::acc::PrivateRecipeOp
 Fortran::lower::createOrGetPrivateRecipe(mlir::OpBuilder &builder,
                                          llvm::StringRef recipeName,
                                          mlir::Location loc, mlir::Type ty) {
   mlir::ModuleOp mod =
       builder.getBlock()->getParent()->getParentOfType<mlir::ModuleOp>();
   if (auto recipe = mod.lookupSymbol<mlir::acc::PrivateRecipeOp>(recipeName))
     return recipe;
 
   auto crtPos = builder.saveInsertionPoint();
   mlir::OpBuilder modBuilder(mod.getBodyRegion());
   auto recipe =
       modBuilder.create<mlir::acc::PrivateRecipeOp>(loc, recipeName, ty);
   builder.createBlock(&recipe.getInitRegion(), recipe.getInitRegion().end(),
                       {ty}, {loc});
   builder.setInsertionPointToEnd(&recipe.getInitRegion().back());
-
-  mlir::Value retVal = recipe.getInitRegion().front().getArgument(0);
-  if (auto refTy = mlir::dyn_cast_or_null<fir::ReferenceType>(ty)) {
-    if (fir::isa_trivial(refTy.getEleTy()))
-      retVal = builder.create<fir::AllocaOp>(loc, refTy.getEleTy());
-    else if (auto seqTy =
-                 mlir::dyn_cast_or_null<fir::SequenceType>(refTy.getEleTy())) {
-      if (seqTy.hasDynamicExtents())
-        TODO(loc, "private recipe of array with dynamic extents");
-      if (fir::isa_trivial(seqTy.getEleTy()))
-        retVal = builder.create<fir::AllocaOp>(loc, seqTy);
-    }
-  }
-  builder.create<mlir::acc::YieldOp>(loc, retVal);
+  genPrivateLikeInitRegion<mlir::acc::PrivateRecipeOp>(builder, recipe, ty,
+                                                       loc);
   builder.restoreInsertionPoint(crtPos);
   return recipe;
 }
 
 mlir::acc::FirstprivateRecipeOp Fortran::lower::createOrGetFirstprivateRecipe(
     mlir::OpBuilder &builder, llvm::StringRef recipeName, mlir::Location loc,
     mlir::Type ty) {
   mlir::ModuleOp mod =
       builder.getBlock()->getParent()->getParentOfType<mlir::ModuleOp>();
   if (auto recipe =
           mod.lookupSymbol<mlir::acc::FirstprivateRecipeOp>(recipeName))
     return recipe;
 
   auto crtPos = builder.saveInsertionPoint();
   mlir::OpBuilder modBuilder(mod.getBodyRegion());
   auto recipe =
       modBuilder.create<mlir::acc::FirstprivateRecipeOp>(loc, recipeName, ty);
   builder.createBlock(&recipe.getInitRegion(), recipe.getInitRegion().end(),
                       {ty}, {loc});
   builder.setInsertionPointToEnd(&recipe.getInitRegion().back());
-  builder.create<mlir::acc::YieldOp>(
-      loc, recipe.getInitRegion().front().getArgument(0));
+  genPrivateLikeInitRegion<mlir::acc::FirstprivateRecipeOp>(builder, recipe, ty,
+                                                            loc);
 
   // Add empty copy region for firstprivate. TODO add copy sequence.
   builder.createBlock(&recipe.getCopyRegion(), recipe.getCopyRegion().end(),
                       {ty, ty}, {loc, loc});
+
   builder.setInsertionPointToEnd(&recipe.getCopyRegion().back());
+  if (auto refTy = mlir::dyn_cast_or_null<fir::ReferenceType>(ty)) {
+    if (fir::isa_trivial(refTy.getEleTy())) {
+      mlir::Value initValue = builder.create<fir::LoadOp>(
+          loc, recipe.getCopyRegion().front().getArgument(0));
+      builder.create<fir::StoreOp>(
+          loc, initValue, recipe.getCopyRegion().front().getArgument(1));
+    } else if (auto seqTy = mlir::dyn_cast_or_null<fir::SequenceType>(
+                   refTy.getEleTy())) {
+      if (seqTy.hasDynamicExtents())
+        TODO(loc, "private recipe of array with dynamic extents");
+      mlir::Type idxTy = builder.getIndexType();
+      mlir::Type refTy = fir::ReferenceType::get(seqTy.getEleTy());
+      mlir::Value arraySrc = recipe.getCopyRegion().front().getArgument(0);
+      mlir::Value arrayDst = recipe.getCopyRegion().front().getArgument(1);
+      llvm::SmallVector<fir::DoLoopOp> loops;
+      llvm::SmallVector<mlir::Value> ivs;
+      for (auto ext : llvm::reverse(seqTy.getShape())) {
+        auto lb = builder.create<mlir::arith::ConstantOp>(
+            loc, idxTy, builder.getIntegerAttr(idxTy, 0));
+        auto ub = builder.create<mlir::arith::ConstantOp>(
+            loc, idxTy, builder.getIntegerAttr(idxTy, ext - 1));
+        auto step = builder.create<mlir::arith::ConstantOp>(
+            loc, idxTy, builder.getIntegerAttr(idxTy, 1));
+        auto loop = builder.create<fir::DoLoopOp>(loc, lb, ub, step,
+                                                  /*unordered=*/false);
+        builder.setInsertionPointToStart(loop.getBody());
+        loops.push_back(loop);
+        ivs.push_back(loop.getInductionVar());
+      }
+      auto addr1 = builder.create<fir::CoordinateOp>(loc, refTy, arraySrc, ivs);
+      auto addr2 = builder.create<fir::CoordinateOp>(loc, refTy, arrayDst, ivs);
+      auto loadedValue = builder.create<fir::LoadOp>(loc, addr1);
+      builder.create<fir::StoreOp>(loc, loadedValue, addr2);
+      builder.setInsertionPointAfter(loops[0]);
+    }
+  }
   builder.create<mlir::acc::TerminatorOp>(loc);
-
   builder.restoreInsertionPoint(crtPos);
   return recipe;
 }
 
 /// Rebuild the array type from the acc.bounds operation with constant
 /// lowerbound/upperbound or extent.
 mlir::Type getTypeFromBounds(llvm::SmallVector<mlir::Value> &bounds,
                              mlir::Type ty) {

flang/test/Lower/OpenACC/acc-parallel-loop.f90

 ! This test checks lowering of OpenACC parallel loop combined directive.
 
 ! RUN: bbc -fopenacc -emit-fir %s -o - | FileCheck %s
 
 ! CHECK-LABEL: acc.firstprivate.recipe @firstprivatization_ref_10xf32 : !fir.ref<!fir.array<10xf32>> init {
-! CHECK: ^bb0(%arg0: !fir.ref<!fir.array<10xf32>>):
-! CHECK:   acc.yield %arg0 : !fir.ref<!fir.array<10xf32>>
+! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<10xf32>>):
+! CHECK:   %[[ALLOCA:.*]] = fir.alloca !fir.array<10xf32> 
+! CHECK:   acc.yield %[[ALLOCA]] : !fir.ref<!fir.array<10xf32>>
 ! CHECK: } copy {
-! CHECK:  ^bb0(%arg0: !fir.ref<!fir.array<10xf32>>, %arg1: !fir.ref<!fir.array<10xf32>>):
+! CHECK: ^bb0(%[[SRC:.*]]: !fir.ref<!fir.array<10xf32>>, %[[DST:.*]]: !fir.ref<!fir.array<10xf32>>):
+! CHECK:   %[[LB0:.*]] = arith.constant 0 : index
+! CHECK:   %[[UB0:.*]] = arith.constant 9 : index
+! CHECK:   %[[STEP0:.*]] = arith.constant 1 : index
+! CHECK:   fir.do_loop %[[IV0:.*]] = %[[LB0]] to %[[UB0]] step %[[STEP0]] {
+! CHECK:     %[[COORD0:.*]] = fir.coordinate_of %[[SRC]], %[[IV0]] : (!fir.ref<!fir.array<10xf32>>, index) -> !fir.ref<f32>
+! CHECK:     %[[COORD1:.*]] = fir.coordinate_of %[[DST]], %[[IV0]] : (!fir.ref<!fir.array<10xf32>>, index) -> !fir.ref<f32>
+! CHECK:     %[[LOAD:.*]] = fir.load %[[COORD0]] : !fir.ref<f32>
+! CHECK:     fir.store %[[LOAD]] to %[[COORD1]] : !fir.ref<f32>
+! CHECK:   }
 ! CHECK:   acc.terminator
 ! CHECK: }
 
 ! CHECK-LABEL: acc.private.recipe @privatization_ref_10xf32 : !fir.ref<!fir.array<10xf32>> init {
 ! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<10xf32>>):
 ! CHECK:   acc.yield %{{.*}} : !fir.ref<!fir.array<10xf32>>
 ! CHECK: }
 

flang/test/Lower/OpenACC/acc-parallel.f90

 ! This test checks lowering of OpenACC parallel directive.
 
 ! RUN: bbc -fopenacc -emit-fir %s -o - | FileCheck %s
 
 ! CHECK-LABEL: acc.firstprivate.recipe @firstprivatization_ref_10x10xf32 : !fir.ref<!fir.array<10x10xf32>> init {
-! CHECK: ^bb0(%arg0: !fir.ref<!fir.array<10x10xf32>>):
-! CHECK:   acc.yield %arg0 : !fir.ref<!fir.array<10x10xf32>>
+! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<10x10xf32>>):
+! CHECK:   %[[ALLOCA:.*]] = fir.alloca !fir.array<10x10xf32> 
+! CHECK:   acc.yield %[[ALLOCA]] : !fir.ref<!fir.array<10x10xf32>>
 ! CHECK: } copy {
 ! CHECK: ^bb0(%arg0: !fir.ref<!fir.array<10x10xf32>>, %arg1: !fir.ref<!fir.array<10x10xf32>>):
 ! CHECK:   acc.terminator
 ! CHECK: }
 
 ! CHECK-LABEL: acc.private.recipe @privatization_ref_10x10xf32 : !fir.ref<!fir.array<10x10xf32>> init {
 ! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<10x10xf32>>):
 ! CHECK: acc.yield %{{.*}} : !fir.ref<!fir.array<10x10xf32>>

flang/test/Lower/OpenACC/acc-private.f90

 ! This test checks lowering of OpenACC loop directive.
 
 ! RUN: bbc -fopenacc -emit-fir %s -o - | FileCheck %s
 
+! CHECK-LABEL: acc.firstprivate.recipe @firstprivatization_ref_100xf32 : !fir.ref<!fir.array<100xf32>> init {
+! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<100xf32>>):
+! CHECK:   %[[ALLOCA:.*]] = fir.alloca !fir.array<100xf32>
+! CHECK:   acc.yield %[[ALLOCA]] : !fir.ref<!fir.array<100xf32>>
+! CHECK: } copy {
+! CHECK: ^bb0(%[[SRC:.*]]: !fir.ref<!fir.array<100xf32>>, %[[DST:.*]]: !fir.ref<!fir.array<100xf32>>):
+! CHECK:   %[[LB0:.*]] = arith.constant 0 : index
+! CHECK:   %[[UB0:.*]] = arith.constant 99 : index
+! CHECK:   %[[STEP1:.*]] = arith.constant 1 : index
+! CHECK:   fir.do_loop %[[IV0:.*]] = %c0 to %c99 step %c1 {
+! CHECK:     %[[COORD0:.*]] = fir.coordinate_of %[[SRC]], %[[IV0]] : (!fir.ref<!fir.array<100xf32>>, index) -> !fir.ref<f32>
+! CHECK:     %[[COORD1:.*]] = fir.coordinate_of %[[DST]], %[[IV0]] : (!fir.ref<!fir.array<100xf32>>, index) -> !fir.ref<f32>
+! CHECK:     %[[VALUE:.*]] = fir.load %[[COORD0]] : !fir.ref<f32>
+! CHECK:     fir.store %[[VALUE]] to %[[COORD1]] : !fir.ref<f32>
+! CHECK:   }
+! CHECK:   acc.terminator
+! CHECK: }
+
+! CHECK-LABEL: acc.firstprivate.recipe @firstprivatization_ref_i32 : !fir.ref<i32> init {
+! CHECK: ^bb0(%{{.*}}: !fir.ref<i32>):
+! CHECK: %[[ALLOCA:.*]] = fir.alloca i32
+! CHECK: acc.yield %[[ALLOCA]] : !fir.ref<i32>
+! CHECK: } copy {
+! CHECK: ^bb0(%[[SRC:.*]]: !fir.ref<i32>, %[[DST:.*]]: !fir.ref<i32>):
+! CHECK:   %[[VALUE:.*]] = fir.load %[[SRC]] : !fir.ref<i32>
+! CHECK:   fir.store %[[VALUE]] to %[[DST]] : !fir.ref<i32>
+! CHECK:   acc.terminator
+! CHECK: }
+
 ! CHECK-LABEL: acc.private.recipe @privatization_ref_50xf32 : !fir.ref<!fir.array<50xf32>> init {
 ! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<50xf32>>):
 ! CHECK: %[[ALLOCA:.*]] = fir.alloca !fir.array<50xf32>
 ! CHECK: acc.yield %[[ALLOCA]] : !fir.ref<!fir.array<50xf32>>
 ! CHECK: }
 
 ! CHECK-LABEL: acc.private.recipe @privatization_ref_100xf32 : !fir.ref<!fir.array<100xf32>> init {
 ! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<100xf32>>):

 
 ! CHECK: %[[C1:.*]] = arith.constant 1 : index
 ! CHECK: %[[LB:.*]] = arith.constant 0 : index
 ! CHECK: %[[UB:.*]] = arith.constant 49 : index
 ! CHECK: %[[BOUND:.*]] = acc.bounds lowerbound(%[[LB]] : index) upperbound(%[[UB]] : index) stride(%[[C1]] : index) startIdx(%[[C1]] : index)
 ! CHECK: %[[B_PRIVATE:.*]] = acc.private varPtr(%[[B]] : !fir.ref<!fir.array<100xf32>>) bounds(%[[BOUND]]) -> !fir.ref<!fir.array<50xf32>> {name = "b(1:50)"}
 ! CHECK: acc.loop private(@privatization_ref_50xf32 -> %[[B_PRIVATE]] : !fir.ref<!fir.array<50xf32>>)
 
+  !$acc parallel loop firstprivate(c)
+  DO i = 1, n
+    c = i
+    a(i) = b(i) + c
+  END DO
+  
+! CHECK: %[[FP_C:.*]] = acc.firstprivate varPtr(%[[C]] : !fir.ref<i32>)   -> !fir.ref<i32> {name = "c"}
+! CHECK: acc.parallel firstprivate(@firstprivatization_ref_i32 -> %[[FP_C]] : !fir.ref<i32>)
+! CHECK: acc.yield
+
+  !$acc parallel loop firstprivate(b)
+  DO i = 1, n
+    c = i
+    a(i) = b(i) + c
+  END DO
+
+! CHECK: %[[C1:.*]] = arith.constant 1 : index
+! CHECK: %[[LB:.*]] = arith.constant 0 : index
+! CHECK: %[[UB:.*]] = arith.subi %{{.*}}, %[[C1]] : index
+! CHECK: %[[BOUND:.*]] = acc.bounds lowerbound(%[[LB]] : index) upperbound(%[[UB]] : index) stride(%[[C1]] : index) startIdx(%[[C1]] : index)
+! CHECK: %[[FP_B:.*]] = acc.firstprivate varPtr(%[[B]] : !fir.ref<!fir.array<100xf32>>) bounds(%[[BOUND]]) -> !fir.ref<!fir.array<100xf32>> {name = "b"}
+! CHECK: acc.parallel firstprivate(@firstprivatization_ref_100xf32 -> %[[FP_B]] : !fir.ref<!fir.array<100xf32>>)
 
 end program

flang/test/Lower/OpenACC/acc-serial-loop.f90

 ! This test checks lowering of Openacc serial loop combined directive.
 
 ! RUN: bbc -fopenacc -emit-fir %s -o - | FileCheck %s
 
 ! CHECK-LABEL: acc.firstprivate.recipe @firstprivatization_ref_10xf32 : !fir.ref<!fir.array<10xf32>> init {
-! CHECK: ^bb0(%arg0: !fir.ref<!fir.array<10xf32>>):
-! CHECK:   acc.yield %arg0 : !fir.ref<!fir.array<10xf32>>
+! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<10xf32>>):
+! CHECK:   %[[ALLOCA:.*]] = fir.alloca !fir.array<10xf32>
+! CHECK:   acc.yield %[[ALLOCA]] : !fir.ref<!fir.array<10xf32>>
 ! CHECK: } copy {
 ! CHECK:  ^bb0(%arg0: !fir.ref<!fir.array<10xf32>>, %arg1: !fir.ref<!fir.array<10xf32>>):
 ! CHECK:   acc.terminator
 ! CHECK: }
 
 ! CHECK-LABEL: acc.private.recipe @privatization_ref_10xf32 : !fir.ref<!fir.array<10xf32>> init {
 ! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<10xf32>>):
 ! CHECK:   acc.yield %{{.*}} : !fir.ref<!fir.array<10xf32>>

flang/test/Lower/OpenACC/acc-serial.f90

 ! This test checks lowering of OpenACC serial directive.
 
 ! RUN: bbc -fopenacc -emit-fir %s -o - | FileCheck %s
 
 ! CHECK-LABEL: acc.firstprivate.recipe @firstprivatization_ref_10x10xf32 : !fir.ref<!fir.array<10x10xf32>> init {
-! CHECK: ^bb0(%arg0: !fir.ref<!fir.array<10x10xf32>>):
-! CHECK:   acc.yield %arg0 : !fir.ref<!fir.array<10x10xf32>>
+! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<10x10xf32>>):
+! CHECK:   %[[ALLOCA:.*]] = fir.alloca !fir.array<10x10xf32>
+! CHECK:   acc.yield %[[ALLOCA]] : !fir.ref<!fir.array<10x10xf32>>
 ! CHECK: } copy {
 ! CHECK: ^bb0(%arg0: !fir.ref<!fir.array<10x10xf32>>, %arg1: !fir.ref<!fir.array<10x10xf32>>):
 ! CHECK:   acc.terminator
 ! CHECK: }
 
 ! CHECK-LABEL: acc.private.recipe @privatization_ref_10x10xf32 : !fir.ref<!fir.array<10x10xf32>> init {
 ! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<10x10xf32>>):
 ! CHECK: acc.yield %{{.*}} : !fir.ref<!fir.array<10x10xf32>>