[Mlir-commits] [mlir] [flang][openacc] Support array with dynamic extents in firstprivate recipe (PR #69026)

Valentin Clement バレンタイン クレメン llvmlistbot at llvm.org
Mon Oct 16 10:41:24 PDT 2023


https://github.com/clementval updated https://github.com/llvm/llvm-project/pull/69026

>From 8ace7491928d1ee198bb09399853ddc640f5468b Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Fri, 13 Oct 2023 13:37:52 -0700
Subject: [PATCH 1/2] [flang][openacc] Support array with dynamic extents in
 firstprivate recipe

---
 flang/lib/Lower/OpenACC.cpp                   | 62 +++++++++++--------
 .../test/Lower/OpenACC/acc-parallel-loop.f90  | 27 --------
 flang/test/Lower/OpenACC/acc-private.f90      | 60 ++++++++++++------
 mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp       |  2 +-
 4 files changed, 77 insertions(+), 74 deletions(-)

diff --git a/flang/lib/Lower/OpenACC.cpp b/flang/lib/Lower/OpenACC.cpp
index e09266121cdb997..8a25df7968b784e 100644
--- a/flang/lib/Lower/OpenACC.cpp
+++ b/flang/lib/Lower/OpenACC.cpp
@@ -569,8 +569,20 @@ mlir::acc::FirstprivateRecipeOp Fortran::lower::createOrGetFirstprivateRecipe(
   mlir::OpBuilder modBuilder(mod.getBodyRegion());
   auto recipe =
       modBuilder.create<mlir::acc::FirstprivateRecipeOp>(loc, recipeName, ty);
+  llvm::SmallVector<mlir::Type> initArgsTy{ty};
+  llvm::SmallVector<mlir::Location> initArgsLoc{loc};
+  auto refTy = fir::unwrapRefType(ty);
+  if (auto seqTy = mlir::dyn_cast_or_null<fir::SequenceType>(refTy)) {
+    if (seqTy.hasDynamicExtents()) {
+      mlir::Type idxTy = builder.getIndexType();
+      for (unsigned i = 0; i < seqTy.getDimension(); ++i) {
+        initArgsTy.push_back(idxTy);
+        initArgsLoc.push_back(loc);
+      }
+    }
+  }
   builder.createBlock(&recipe.getInitRegion(), recipe.getInitRegion().end(),
-                      {ty}, {loc});
+                      initArgsTy, initArgsLoc);
   builder.setInsertionPointToEnd(&recipe.getInitRegion().back());
   genPrivateLikeInitRegion<mlir::acc::FirstprivateRecipeOp>(builder, recipe, ty,
                                                             loc);
@@ -601,32 +613,28 @@ mlir::acc::FirstprivateRecipeOp Fortran::lower::createOrGetFirstprivateRecipe(
     builder.create<fir::StoreOp>(loc, initValue,
                                  recipe.getCopyRegion().front().getArgument(1));
   } else if (auto seqTy = mlir::dyn_cast_or_null<fir::SequenceType>(ty)) {
-    if (seqTy.hasDynamicExtents())
-      TODO(loc, "firstprivate 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]);
+    fir::FirOpBuilder firBuilder{builder, recipe.getOperation()};
+    auto shape = genShapeFromBoundsOrArgs(
+        loc, firBuilder, seqTy, bounds, recipe.getCopyRegion().getArguments());
+
+    auto leftDeclOp = builder.create<hlfir::DeclareOp>(
+        loc, recipe.getCopyRegion().getArgument(0), llvm::StringRef{}, shape,
+        llvm::ArrayRef<mlir::Value>{}, fir::FortranVariableFlagsAttr{});
+    auto rightDeclOp = builder.create<hlfir::DeclareOp>(
+        loc, recipe.getCopyRegion().getArgument(1), llvm::StringRef{}, shape,
+        llvm::ArrayRef<mlir::Value>{}, fir::FortranVariableFlagsAttr{});
+
+    hlfir::DesignateOp::Subscripts triplets =
+        getSubscriptsFromArgs(recipe.getCopyRegion().getArguments());
+    auto leftEntity = hlfir::Entity{leftDeclOp.getBase()};
+    auto left =
+        genDesignateWithTriplets(firBuilder, loc, leftEntity, triplets, shape);
+    auto rightEntity = hlfir::Entity{rightDeclOp.getBase()};
+    auto right =
+        genDesignateWithTriplets(firBuilder, loc, rightEntity, triplets, shape);
+
+    firBuilder.create<hlfir::AssignOp>(loc, left, right);
+
   } else if (auto boxTy = mlir::dyn_cast_or_null<fir::BaseBoxType>(ty)) {
     fir::FirOpBuilder firBuilder{builder, recipe.getOperation()};
     llvm::SmallVector<mlir::Value> tripletArgs;
diff --git a/flang/test/Lower/OpenACC/acc-parallel-loop.f90 b/flang/test/Lower/OpenACC/acc-parallel-loop.f90
index 9a27a359e80b73b..80b1272bd1b10b6 100644
--- a/flang/test/Lower/OpenACC/acc-parallel-loop.f90
+++ b/flang/test/Lower/OpenACC/acc-parallel-loop.f90
@@ -3,33 +3,6 @@
 ! RUN: bbc -fopenacc -emit-fir %s -o - | FileCheck %s --check-prefixes=CHECK,FIR
 ! RUN: bbc -fopenacc -emit-hlfir %s -o - | FileCheck %s --check-prefixes=CHECK,HLFIR
 
-! CHECK-LABEL: acc.firstprivate.recipe @firstprivatization_section_ext10_ref_10xf32 : !fir.ref<!fir.array<10xf32>> init {
-! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<10xf32>>):
-! HLFIR:   %[[SHAPE:.*]] = fir.shape %{{.*}} : (index) -> !fir.shape<1>
-! CHECK:   %[[ALLOCA:.*]] = fir.alloca !fir.array<10xf32>
-! HLFIR:   %[[DECLARE:.*]]:2 = hlfir.declare %[[ALLOCA]](%[[SHAPE]]) {uniq_name = "acc.private.init"} : (!fir.ref<!fir.array<10xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xf32>>, !fir.ref<!fir.array<10xf32>>)
-! HLFIR:   acc.yield %[[DECLARE]]#0 : !fir.ref<!fir.array<10xf32>>
-! CHECK: } copy {
-! 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>>):
-! HLFIR:   %[[SHAPE:.*]] = fir.shape %{{.*}} : (index) -> !fir.shape<1>
-! HLFIR:   %[[DECLARE:.*]]:2 = hlfir.declare %[[ALLOCA]](%[[SHAPE]]) {uniq_name = "acc.private.init"} : (!fir.ref<!fir.array<10xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xf32>>, !fir.ref<!fir.array<10xf32>>)
-! HLFIR:   acc.yield %[[DECLARE]]#0 : !fir.ref<!fir.array<10xf32>>
-! CHECK: }
-
 ! CHECK-LABEL: func.func @_QPacc_parallel_loop()
 
 subroutine acc_parallel_loop
diff --git a/flang/test/Lower/OpenACC/acc-private.f90 b/flang/test/Lower/OpenACC/acc-private.f90
index 10c1bfc7c3802a3..9ce1828e63ddf10 100644
--- a/flang/test/Lower/OpenACC/acc-private.f90
+++ b/flang/test/Lower/OpenACC/acc-private.f90
@@ -3,6 +3,23 @@
 ! RUN: bbc -fopenacc -emit-fir %s -o - | FileCheck %s --check-prefixes=CHECK,FIR
 ! RUN: bbc -fopenacc -emit-hlfir %s -o - | FileCheck %s --check-prefixes=CHECK,HLFIR
 
+! CHECK-LABEL: acc.firstprivate.recipe @firstprivatization_ref_UxUx2xi32 : !fir.ref<!fir.array<?x?x2xi32>> init {
+! CHECK: ^bb0(%[[ARG0:.*]]: !fir.ref<!fir.array<?x?x2xi32>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index):
+! HLFIR:   %[[SHAPE:.*]] = fir.shape %[[ARG1]], %[[ARG2]], %[[ARG3]] : (index, index, index) -> !fir.shape<3>
+! HLFIR:   %[[TEMP:.*]] = fir.alloca !fir.array<?x?x2xi32>, %[[ARG1]], %[[ARG2]], %[[ARG3]]
+! HLFIR:   %[[DECL:.*]]:2 = hlfir.declare %[[TEMP]](%[[SHAPE]]) {uniq_name = "acc.private.init"} : (!fir.ref<!fir.array<?x?x2xi32>>, !fir.shape<3>) -> (!fir.box<!fir.array<?x?x2xi32>>, !fir.ref<!fir.array<?x?x2xi32>>)
+! HLFIR:   acc.yield %[[DECL]]#0 : !fir.box<!fir.array<?x?x2xi32>>
+! CHECK: } copy {
+! CHECK: ^bb0(%[[SRC:.*]]: !fir.ref<!fir.array<?x?x2xi32>>, %[[DST:.*]]: !fir.ref<!fir.array<?x?x2xi32>>, %[[LB0:.*]]: index, %[[UB0:.*]]: index, %[[STEP0:.*]]: index, %[[LB1:.*]]: index, %[[UB1:.*]]: index, %[[STEP1:.*]]: index, %[[LB2:.*]]: index, %[[UB2:.*]]: index, %[[STEP2:.*]]: index):
+! HLFIR:   %[[SHAPE:.*]] = fir.shape %{{.*}}, %{{.*}}, %{{.*}} : (index, index, index) -> !fir.shape<3>
+! HLFIR:   %[[DECL_SRC:.*]]:2 = hlfir.declare %[[SRC]](%[[SHAPE]]) {uniq_name = ""} : (!fir.ref<!fir.array<?x?x2xi32>>, !fir.shape<3>) -> (!fir.box<!fir.array<?x?x2xi32>>, !fir.ref<!fir.array<?x?x2xi32>>)
+! HLFIR:   %[[DECL_DST:.*]]:2 = hlfir.declare %[[DST]](%[[SHAPE]]) {uniq_name = ""} : (!fir.ref<!fir.array<?x?x2xi32>>, !fir.shape<3>) -> (!fir.box<!fir.array<?x?x2xi32>>, !fir.ref<!fir.array<?x?x2xi32>>)
+! HLFIR:   %[[DES_SRC:.*]] = hlfir.designate %[[DECL_SRC]]#0 (%[[LB0]]:%[[UB0]]:%[[STEP0]], %[[LB1]]:%[[UB1]]:%[[STEP1]], %[[LB2]]:%[[UB2]]:%[[STEP2]]) shape %[[SHAPE]] : (!fir.box<!fir.array<?x?x2xi32>>, index, index, index, index, index, index, index, index, index, !fir.shape<3>) -> !fir.box<!fir.array<?x?x2xi32>>
+! HLFIR:   %[[DES_DST:.*]] = hlfir.designate %[[DECL_DST]]#0 (%[[LB0]]:%[[UB0]]:%[[STEP0]], %[[LB1]]:%[[UB1]]:%[[STEP1]], %[[LB2]]:%[[UB2]]:%[[STEP2]]) shape %[[SHAPE]] : (!fir.box<!fir.array<?x?x2xi32>>, index, index, index, index, index, index, index, index, index, !fir.shape<3>) -> !fir.box<!fir.array<?x?x2xi32>>
+! HLFIR:   hlfir.assign %[[DES_SRC]] to %[[DES_DST]] : !fir.box<!fir.array<?x?x2xi32>>, !fir.box<!fir.array<?x?x2xi32>>
+! HLFIR:   acc.terminator
+! CHECK: }
+
 ! CHECK-LABEL: acc.firstprivate.recipe @firstprivatization_section_lb4.ub9_box_Uxi32 : !fir.box<!fir.array<?xi32>> init {
 ! CHECK: ^bb0(%{{.*}}: !fir.box<!fir.array<?xi32>>):
 ! CHECK: } copy {
@@ -87,16 +104,12 @@
 ! HLFIR:   acc.yield %[[DECLARE]]#0 : !fir.ref<!fir.array<50xf32>>
 ! CHECK: } copy {
 ! CHECK: ^bb0(%[[SRC:.*]]: !fir.ref<!fir.array<50xf32>>, %[[DST:.*]]: !fir.ref<!fir.array<50xf32>>):
-! CHECK:   %[[LB0:.*]] = arith.constant 0 : index
-! CHECK:   %[[UB0:.*]] = arith.constant 49 : 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<50xf32>>, index) -> !fir.ref<f32>
-! CHECK:     %[[COORD1:.*]] = fir.coordinate_of %[[DST]], %[[IV0]] : (!fir.ref<!fir.array<50xf32>>, 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
+! HLFIR:   %[[SHAPE:.*]] = fir.shape %{{.*}} : (index) -> !fir.shape<1>
+! HLFIR:   %[[DECL_SRC:.*]]:2 = hlfir.declare %[[SRC]](%[[SHAPE]]) {uniq_name = ""} : (!fir.ref<!fir.array<50xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<50xf32>>, !fir.ref<!fir.array<50xf32>>)
+! HLFIR:   %[[DECL_DST:.*]]:2 = hlfir.declare %[[DST]](%[[SHAPE]]) {uniq_name = ""} : (!fir.ref<!fir.array<50xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<50xf32>>, !fir.ref<!fir.array<50xf32>>)
+! HLFIR:   %[[DES_SRC:.*]] = hlfir.designate %[[DECL_SRC]]#0 shape %[[SHAPE:.*]] : (!fir.ref<!fir.array<50xf32>>, !fir.shape<1>) -> !fir.ref<!fir.array<50xf32>>
+! HLFIR:   %[[DES_DST:.*]] = hlfir.designate %[[DECL_DST]]#0 shape %[[SHAPE:.*]] : (!fir.ref<!fir.array<50xf32>>, !fir.shape<1>) -> !fir.ref<!fir.array<50xf32>>
+! HLFIR:   hlfir.assign %[[DES_SRC]] to %[[DES_DST]] : !fir.ref<!fir.array<50xf32>>, !fir.ref<!fir.array<50xf32>>
 ! CHECK: }
 
 ! CHECK-LABEL: acc.firstprivate.recipe @firstprivatization_section_ext100_ref_100xf32 : !fir.ref<!fir.array<100xf32>> init {
@@ -107,15 +120,12 @@
 ! HLFIR:   acc.yield %[[DECLARE]]#0 : !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:   }
+! HLFIR:   %[[SHAPE:.*]] = fir.shape %{{.*}} : (index) -> !fir.shape<1>
+! HLFIR:   %[[DECL_SRC:.*]]:2 = hlfir.declare %[[SRC]](%[[SHAPE]]) {uniq_name = ""} : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<100xf32>>, !fir.ref<!fir.array<100xf32>>)
+! HLFIR:   %[[DECL_DST:.*]]:2 = hlfir.declare %[[DST]](%[[SHAPE]]) {uniq_name = ""} : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<100xf32>>, !fir.ref<!fir.array<100xf32>>)
+! HLFIR:   %[[DES_SRC:.*]] = hlfir.designate %[[DECL_SRC]]#0 shape %[[SHAPE]] : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>) -> !fir.ref<!fir.array<100xf32>>
+! HLFIR:   %[[DES_DST:.*]] = hlfir.designate %[[DECL_DST]]#0 shape %[[SHAPE]] : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>) -> !fir.ref<!fir.array<100xf32>>
+! HLFIR:   hlfir.assign %[[DES_SRC]] to %[[DES_DST]] : !fir.ref<!fir.array<100xf32>>, !fir.ref<!fir.array<100xf32>>
 ! CHECK:   acc.terminator
 ! CHECK: }
 
@@ -337,3 +347,15 @@ subroutine acc_firstprivate_assumed_shape_with_section(a, n)
     a(i) = i
   end do
 end subroutine
+
+subroutine acc_firstprivate_dynamic_extent(a, n)
+  integer :: n, i
+  integer :: a(n, n, 2)
+
+  !$acc parallel loop firstprivate(a)
+  do i = 1, n
+    a(i, i, 1) = i
+  end do
+end subroutine
+
+! CHECK: acc.parallel firstprivate(@firstprivatization_ref_UxUx2xi32 -> %{{.*}} : !fir.ref<!fir.array<?x?x2xi32>>)
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
index cea93b8a2ca8ceb..b7e2aec6a4e6a92 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
@@ -452,7 +452,7 @@ LogicalResult acc::PrivateRecipeOp::verifyRegions() {
 LogicalResult acc::FirstprivateRecipeOp::verifyRegions() {
   if (failed(verifyInitLikeSingleArgRegion(*this, getInitRegion(),
                                            "privatization", "init", getType(),
-                                           /*verifyYield=*/true)))
+                                           /*verifyYield=*/false)))
     return failure();
 
   if (getCopyRegion().empty())

>From 660ef8354ce5faca1488c5f7e77f38e019a1240f Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Fri, 13 Oct 2023 14:23:06 -0700
Subject: [PATCH 2/2] Fix OpenACC mlir test

---
 mlir/test/Dialect/OpenACC/invalid.mlir | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/mlir/test/Dialect/OpenACC/invalid.mlir b/mlir/test/Dialect/OpenACC/invalid.mlir
index 225a8766fc550f1..ff92eab478bb4f5 100644
--- a/mlir/test/Dialect/OpenACC/invalid.mlir
+++ b/mlir/test/Dialect/OpenACC/invalid.mlir
@@ -312,18 +312,6 @@ acc.firstprivate.recipe @privatization_i32 : !llvm.ptr<i32> init {
 
 // -----
 
-// expected-error at +1 {{expects init region to yield a value of the privatization type}}
-acc.firstprivate.recipe @privatization_i32 : !llvm.ptr<f32> init {
-^bb0(%arg0 : !llvm.ptr<f32>):
-  %c1 = arith.constant 1 : i32
-  %c0 = arith.constant 0 : i32
-  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<i32>
-  llvm.store %c0, %0 : !llvm.ptr<i32>
-  acc.yield %0 : !llvm.ptr<i32>
-} copy {}
-
-// -----
-
 // expected-error at +1 {{expects non-empty copy region}}
 acc.firstprivate.recipe @privatization_i32 : !llvm.ptr<i32> init {
 ^bb0(%arg0 : !llvm.ptr<i32>):



More information about the Mlir-commits mailing list