[Flang][OpenMP] Update DO CONCURRENT conversion for the device

This patch makes changes to the `DoConcurrentConversion` pass to follow the "hoisted omp.parallel" representation when converting `do concurrent` constructs into `target teams distribute parallel do`.
ROCm · Aug 13, 2024 · 9a95922 · 9a95922
1 parent 01c368d
commit 9a95922
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 22 deletions.
diff --git a/flang/lib/Optimizer/Transforms/DoConcurrentConversion.cpp b/flang/lib/Optimizer/Transforms/DoConcurrentConversion.cpp
@@ -599,9 +599,7 @@ class DoConcurrentConversion : public mlir::OpConversionPattern<fir::DoLoopOp> {
 
  targetOp = genTargetOp(doLoop.getLoc(), rewriter, mapper,
  outermostLoopLives, targetClauseOps);
- genTeamsOp(doLoop.getLoc(), rewriter, loopNest, mapper,
- loopNestClauseOps);
- genDistributeOp(doLoop.getLoc(), rewriter);
+ genTeamsOp(doLoop.getLoc(), rewriter);
  }
 
  mlir::omp::ParallelOp parallelOp = genParallelOp(
@@ -611,6 +609,9 @@ class DoConcurrentConversion : public mlir::OpConversionPattern<fir::DoLoopOp> {
  looputils::localizeLoopLocalValue(local, parallelOp.getRegion(),
  rewriter);
 
+ if (mapToDevice)
+ genDistributeOp(doLoop.getLoc(), rewriter);
+
  mlir::omp::LoopNestOp ompLoopNest =
  genWsLoopOp(rewriter, loopNest.back().first, mapper, loopNestClauseOps);
 
@@ -800,18 +801,14 @@ class DoConcurrentConversion : public mlir::OpConversionPattern<fir::DoLoopOp> {
  }
 
  mlir::omp::TeamsOp
- genTeamsOp(mlir::Location loc, mlir::ConversionPatternRewriter &rewriter,
- looputils::LoopNestToIndVarMap &loopNest, mlir::IRMapping &mapper,
- mlir::omp::LoopNestOperands &loopNestClauseOps) const {
+ genTeamsOp(mlir::Location loc,
+ mlir::ConversionPatternRewriter &rewriter) const {
  auto teamsOp = rewriter.create<mlir::omp::TeamsOp>(
  loc, /*clauses=*/mlir::omp::TeamsOperands{});
 
  rewriter.createBlock(&teamsOp.getRegion());
  rewriter.setInsertionPoint(rewriter.create<mlir::omp::TerminatorOp>(loc));
 
- genLoopNestIndVarAllocs(rewriter, loopNest, mapper);
- genLoopNestClauseOps(loc, rewriter, loopNest, mapper, loopNestClauseOps);
-
  return teamsOp;
  }
 
@@ -905,12 +902,8 @@ class DoConcurrentConversion : public mlir::OpConversionPattern<fir::DoLoopOp> {
  rewriter.createBlock(&parallelOp.getRegion());
  rewriter.setInsertionPoint(rewriter.create<mlir::omp::TerminatorOp>(loc));
 
- // If mapping to host, the local induction variable and loop bounds need to
- // be emitted as part of the `omp.parallel` op.
- if (!mapToDevice) {
- genLoopNestIndVarAllocs(rewriter, loopNest, mapper);
- genLoopNestClauseOps(loc, rewriter, loopNest, mapper, loopNestClauseOps);
- }
+ genLoopNestIndVarAllocs(rewriter, loopNest, mapper);
+ genLoopNestClauseOps(loc, rewriter, loopNest, mapper, loopNestClauseOps);
 
  return parallelOp;
  }

diff --git a/flang/test/Transforms/DoConcurrent/basic_device.f90 b/flang/test/Transforms/DoConcurrent/basic_device.f90
@@ -43,6 +43,7 @@ program do_concurrent_basic
 
  ! CHECK: %[[A_DEV_DECL:.*]]:2 = hlfir.declare %[[A_ARG]]
  ! CHECK: omp.teams {
+ ! CHECK-NEXT: omp.parallel {
 
  ! CHECK-NEXT: %[[ITER_VAR:.*]] = fir.alloca i32 {bindc_name = "i"}
  ! CHECK-NEXT: %[[BINDING:.*]]:2 = hlfir.declare %[[ITER_VAR]] {uniq_name = "_QFEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
@@ -54,8 +55,6 @@ program do_concurrent_basic
  ! CHECK: %[[STEP:.*]] = arith.constant 1 : index
 
  ! CHECK-NEXT: omp.distribute {
- ! CHECK-NEXT: omp.parallel {
-
  ! CHECK-NEXT: omp.wsloop {
 
  ! CHECK-NEXT: omp.loop_nest (%[[ARG0:.*]]) : index = (%[[LB]]) to (%[[UB]]) inclusive step (%[[STEP]]) {

diff --git a/flang/test/Transforms/DoConcurrent/multiple_iteration_ranges.f90 b/flang/test/Transforms/DoConcurrent/multiple_iteration_ranges.f90
@@ -65,7 +65,7 @@ program main
 ! DEVICE: omp.target
 ! DEVICE: omp.teams
 
-! HOST: omp.parallel {
+! COMMON: omp.parallel {
 
 ! COMMON-NEXT: %[[ITER_VAR_I:.*]] = fir.alloca i32 {bindc_name = "i"}
 ! COMMON-NEXT: %[[BINDING_I:.*]]:2 = hlfir.declare %[[ITER_VAR_I]] {uniq_name = "_QFEi"}
@@ -94,8 +94,7 @@ program main
 ! COMMON: %[[UB_K:.*]] = fir.convert %[[C30]] : (i32) -> index
 ! COMMON: %[[STEP_K:.*]] = arith.constant 1 : index
 
-! DEVICE: omp.distribute
-! DEVICE-NEXT: omp.parallel
+! DEVICE: omp.distribute
 
 ! COMMON: omp.wsloop {
 ! COMMON-NEXT: omp.loop_nest

diff --git a/flang/test/Transforms/DoConcurrent/not_perfectly_nested.f90 b/flang/test/Transforms/DoConcurrent/not_perfectly_nested.f90
@@ -39,9 +39,11 @@ program main
 ! DEVICE: %[[TARGET_K_DECL:.*]]:2 = hlfir.declare %[[K_ARG]] {uniq_name = "_QFEk"}
 
 ! DEVICE: omp.teams
-! DEVICE: omp.distribute
 
 ! COMMON: omp.parallel {
+
+! DEVICE: omp.distribute
+
 ! COMMON: omp.wsloop {
 ! COMMON: omp.loop_nest ({{[^[:space:]]+}}) {{.*}} {
 ! COMMON: fir.do_loop %[[J_IV:.*]] = {{.*}} {

diff --git a/flang/test/Transforms/DoConcurrent/skip_all_nested_loops.f90 b/flang/test/Transforms/DoConcurrent/skip_all_nested_loops.f90
@@ -39,9 +39,11 @@ program main
 ! DEVICE: %[[TARGET_K_DECL:.*]]:2 = hlfir.declare %[[K_ARG]] {uniq_name = "_QFEk"}
 
 ! DEVICE: omp.teams
-! DEVICE: omp.distribute
 
 ! COMMON: omp.parallel {
+
+! DEVICE: omp.distribute
+
 ! COMMON: omp.wsloop {
 ! COMMON: omp.loop_nest ({{[^[:space:]]+}}) {{.*}} {
 ! COMMON: fir.do_loop {{.*}} iter_args(%[[J_IV:.*]] = {{.*}}) -> {{.*}} {