[MLIR][Flang][OpenMP] Host evaluation of trip count for Generic-SPMD

This patch enables the initialization of the loop trip count for kernels representing `target teams distribute` or equivalent constructs. This involves updates to Flang lowering to make sure loop bounds are lowered in advance, omp.target verifier changes and tweaks to MLIR to LLVM IR translation.
ROCm · Nov 26, 2024 · 696d92b · 696d92b
1 parent 67697c6
commit 696d92b
Show file tree

Hide file tree

Showing 11 changed files with 395 additions and 52 deletions.
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -578,7 +578,6 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
     HostEvalInfo &hostInfo = hostEvalInfo.back();
 
     switch (extractOmpDirective(*ompEval)) {
-    // Cases where 'teams' and target SPMD clauses might be present.
     case OMPD_teams_distribute_parallel_do:
     case OMPD_teams_distribute_parallel_do_simd:
       cp.processThreadLimit(stmtCtx, hostInfo.ops);
@@ -589,39 +588,38 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
       [[fallthrough]];
     case OMPD_distribute_parallel_do:
     case OMPD_distribute_parallel_do_simd:
-      cp.processCollapse(loc, eval, hostInfo.ops, hostInfo.iv);
       cp.processNumThreads(stmtCtx, hostInfo.ops);
+      [[fallthrough]];
+    case OMPD_distribute:
+    case OMPD_distribute_simd:
+      cp.processCollapse(loc, eval, hostInfo.ops, hostInfo.iv);
       break;
 
-    // Cases where 'teams' clauses might be present, and target SPMD is
-    // possible by looking at nested evaluations.
     case OMPD_teams:
       cp.processThreadLimit(stmtCtx, hostInfo.ops);
       [[fallthrough]];
     case OMPD_target_teams:
       cp.processNumTeams(stmtCtx, hostInfo.ops);
-      processSingleNestedIf([](Directive nestedDir) {
-        return nestedDir == OMPD_distribute_parallel_do ||
-               nestedDir == OMPD_distribute_parallel_do_simd;
-      });
+      processSingleNestedIf(
+          [](Directive nestedDir) { return topDistributeSet.test(nestedDir); });
       break;
 
-    // Cases where only 'teams' host-evaluated clauses might be present.
     case OMPD_teams_distribute:
     case OMPD_teams_distribute_simd:
       cp.processThreadLimit(stmtCtx, hostInfo.ops);
       [[fallthrough]];
     case OMPD_target_teams_distribute:
     case OMPD_target_teams_distribute_simd:
+      cp.processCollapse(loc, eval, hostInfo.ops, hostInfo.iv);
       cp.processNumTeams(stmtCtx, hostInfo.ops);
       break;
 
     // Standalone 'target' case.
-    case OMPD_target: {
+    case OMPD_target:
       processSingleNestedIf(
           [](Directive nestedDir) { return topTeamsSet.test(nestedDir); });
       break;
-    }
+
     default:
       break;
     }

diff --git a/flang/test/Lower/OpenMP/eval-outside-target.f90 b/flang/test/Lower/OpenMP/eval-outside-target.f90
@@ -155,3 +155,106 @@ subroutine distribute_parallel_do_simd()
   !$omp end distribute parallel do simd
   !$omp end teams
 end subroutine distribute_parallel_do_simd
+
+! BOTH-LABEL: func.func @_QPdistribute
+subroutine distribute()
+  ! BOTH: omp.target
+
+  ! HOST-SAME: host_eval(%{{.*}} -> %[[LB:.*]], %{{.*}} -> %[[UB:.*]], %{{.*}} -> %[[STEP:.*]] : i32, i32, i32)
+
+  ! DEVICE-NOT: host_eval({{.*}})
+  ! DEVICE-SAME: {
+
+  ! BOTH: omp.teams
+  !$omp target teams
+
+  ! BOTH: omp.distribute
+  ! BOTH-NEXT: omp.loop_nest
+
+  ! HOST-SAME: (%{{.*}}) : i32 = (%[[LB]]) to (%[[UB]]) inclusive step (%[[STEP]])
+  !$omp distribute
+  do i=1,10
+    call foo()
+  end do
+  !$omp end distribute
+  !$omp end target teams
+
+  ! BOTH: omp.target
+  ! BOTH-NOT: host_eval({{.*}})
+  ! BOTH-SAME: {
+  ! BOTH: omp.teams
+  !$omp target teams
+  call foo() !< Prevents this from being Generic-SPMD.
+
+  ! BOTH: omp.distribute
+  !$omp distribute
+  do i=1,10
+    call foo()
+  end do
+  !$omp end distribute
+  !$omp end target teams
+
+  ! BOTH: omp.teams
+  !$omp teams
+
+  ! BOTH: omp.distribute
+  !$omp distribute
+  do i=1,10
+    call foo()
+  end do
+  !$omp end distribute
+  !$omp end teams
+end subroutine distribute
+
+! BOTH-LABEL: func.func @_QPdistribute_simd
+subroutine distribute_simd()
+  ! BOTH: omp.target
+
+  ! HOST-SAME: host_eval(%{{.*}} -> %[[LB:.*]], %{{.*}} -> %[[UB:.*]], %{{.*}} -> %[[STEP:.*]] : i32, i32, i32)
+
+  ! DEVICE-NOT: host_eval({{.*}})
+  ! DEVICE-SAME: {
+
+  ! BOTH: omp.teams
+  !$omp target teams
+
+  ! BOTH: omp.distribute
+  ! BOTH-NEXT: omp.simd
+  ! BOTH-NEXT: omp.loop_nest
+
+  ! HOST-SAME: (%{{.*}}) : i32 = (%[[LB]]) to (%[[UB]]) inclusive step (%[[STEP]])
+  !$omp distribute simd
+  do i=1,10
+    call foo()
+  end do
+  !$omp end distribute simd
+  !$omp end target teams
+
+  ! BOTH: omp.target
+  ! BOTH-NOT: host_eval({{.*}})
+  ! BOTH-SAME: {
+  ! BOTH: omp.teams
+  !$omp target teams
+  call foo() !< Prevents this from being Generic-SPMD.
+
+  ! BOTH: omp.distribute
+  ! BOTH-NEXT: omp.simd
+  !$omp distribute simd
+  do i=1,10
+    call foo()
+  end do
+  !$omp end distribute simd
+  !$omp end target teams
+
+  ! BOTH: omp.teams
+  !$omp teams
+
+  ! BOTH: omp.distribute
+  ! BOTH-NEXT: omp.simd
+  !$omp distribute simd
+  do i=1,10
+    call foo()
+  end do
+  !$omp end distribute simd
+  !$omp end teams
+end subroutine distribute_simd
diff --git a/flang/test/Lower/OpenMP/target-generic-spmd.f90 b/flang/test/Lower/OpenMP/target-generic-spmd.f90
@@ -0,0 +1,191 @@
+! RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s
+
+! CHECK-LABEL: func.func @_QPdistribute_generic() {
+subroutine distribute_generic()
+  ! CHECK: omp.target
+  ! CHECK-NOT: host_eval({{.*}})
+  ! CHECK-SAME: {
+  !$omp target
+  !$omp teams
+  !$omp distribute
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute
+  call bar() !< Prevents this from being Generic-SPMD.
+  !$omp end teams
+  !$omp end target
+
+  ! CHECK: omp.target
+  ! CHECK-NOT: host_eval({{.*}})
+  ! CHECK-SAME: {
+  !$omp target teams
+  !$omp distribute
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute
+  call bar() !< Prevents this from being Generic-SPMD.
+  !$omp end target teams
+
+  ! CHECK: omp.target
+  ! CHECK-NOT: host_eval({{.*}})
+  ! CHECK-SAME: {
+  !$omp target teams
+  !$omp distribute
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute
+
+  !$omp distribute
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute
+  !$omp end target teams
+end subroutine distribute_generic
+
+! CHECK-LABEL: func.func @_QPdistribute_spmd() {
+subroutine distribute_spmd()
+  ! CHECK: omp.target
+  ! CHECK-SAME: host_eval({{.*}})
+  !$omp target
+  !$omp teams
+  !$omp distribute
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute
+  !$omp end teams
+  !$omp end target
+
+  ! CHECK: omp.target
+  ! CHECK-SAME: host_eval({{.*}})
+  !$omp target teams
+  !$omp distribute
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute
+  !$omp end target teams
+end subroutine distribute_spmd
+
+! CHECK-LABEL: func.func @_QPdistribute_simd_generic() {
+subroutine distribute_simd_generic()
+  ! CHECK: omp.target
+  ! CHECK-NOT: host_eval({{.*}})
+  ! CHECK-SAME: {
+  !$omp target
+  !$omp teams
+  !$omp distribute simd
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute simd
+  call bar() !< Prevents this from being Generic-SPMD.
+  !$omp end teams
+  !$omp end target
+
+  ! CHECK: omp.target
+  ! CHECK-NOT: host_eval({{.*}})
+  ! CHECK-SAME: {
+  !$omp target teams
+  !$omp distribute simd
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute simd
+  call bar() !< Prevents this from being Generic-SPMD.
+  !$omp end target teams
+
+  ! CHECK: omp.target
+  ! CHECK-NOT: host_eval({{.*}})
+  ! CHECK-SAME: {
+  !$omp target teams
+  !$omp distribute simd
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute simd
+
+  !$omp distribute simd
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute simd
+  !$omp end target teams
+end subroutine distribute_simd_generic
+
+! CHECK-LABEL: func.func @_QPdistribute_simd_spmd() {
+subroutine distribute_simd_spmd()
+  ! CHECK: omp.target
+  ! CHECK-SAME: host_eval({{.*}})
+  !$omp target
+  !$omp teams
+  !$omp distribute simd
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute simd
+  !$omp end teams
+  !$omp end target
+
+  ! CHECK: omp.target
+  ! CHECK-SAME: host_eval({{.*}})
+  !$omp target teams
+  !$omp distribute simd
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end distribute simd
+  !$omp end target teams
+end subroutine distribute_simd_spmd
+
+! CHECK-LABEL: func.func @_QPteams_distribute_spmd() {
+subroutine teams_distribute_spmd()
+  ! CHECK: omp.target
+  ! CHECK-SAME: host_eval({{.*}})
+  !$omp target
+  !$omp teams distribute
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end teams distribute
+  !$omp end target
+end subroutine teams_distribute_spmd
+
+! CHECK-LABEL: func.func @_QPteams_distribute_simd_spmd() {
+subroutine teams_distribute_simd_spmd()
+  ! CHECK: omp.target
+  ! CHECK-SAME: host_eval({{.*}})
+  !$omp target
+  !$omp teams distribute simd
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end teams distribute simd
+  !$omp end target
+end subroutine teams_distribute_simd_spmd
+
+! CHECK-LABEL: func.func @_QPtarget_teams_distribute_spmd() {
+subroutine target_teams_distribute_spmd()
+  ! CHECK: omp.target
+  ! CHECK-SAME: host_eval({{.*}})
+  !$omp target teams distribute
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end target teams distribute
+end subroutine target_teams_distribute_spmd
+
+! CHECK-LABEL: func.func @_QPtarget_teams_distribute_simd_spmd() {
+subroutine target_teams_distribute_simd_spmd()
+  ! CHECK: omp.target
+  ! CHECK-SAME: host_eval({{.*}})
+  !$omp target teams distribute simd
+  do i = 1, 10
+    call foo(i)
+  end do
+  !$omp end target teams distribute simd
+end subroutine target_teams_distribute_simd_spmd
diff --git a/flang/test/Lower/OpenMP/target-spmd.f90 b/flang/test/Lower/OpenMP/target-spmd.f90
@@ -12,7 +12,7 @@ subroutine distribute_parallel_do_generic()
     call foo(i)
   end do
   !$omp end distribute parallel do
-  call bar()
+  call bar() !< Prevents this from being SPMD.
   !$omp end teams
   !$omp end target
 
@@ -25,7 +25,7 @@ subroutine distribute_parallel_do_generic()
     call foo(i)
   end do
   !$omp end distribute parallel do
-  call bar()
+  call bar() !< Prevents this from being SPMD.
   !$omp end target teams
 
   ! CHECK: omp.target
@@ -83,7 +83,7 @@ subroutine distribute_parallel_do_simd_generic()
     call foo(i)
   end do
   !$omp end distribute parallel do simd
-  call bar()
+  call bar() !< Prevents this from being SPMD.
   !$omp end teams
   !$omp end target
 
@@ -96,7 +96,7 @@ subroutine distribute_parallel_do_simd_generic()
     call foo(i)
   end do
   !$omp end distribute parallel do simd
-  call bar()
+  call bar() !< Prevents this from being SPMD.
   !$omp end target teams
 
   ! CHECK: omp.target