[flang][OpenMP] Support host_eval for target teams loop

Extends `host_eval` support for the currently supported form of the generic `loop` directive.
ROCm · Dec 11, 2024 · 6b6cf62 · 6b6cf62
1 parent 73deda0
commit 6b6cf62
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 8 deletions.
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -557,8 +557,9 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
       [[fallthrough]];
     case OMPD_target_teams:
       cp.processNumTeams(stmtCtx, hostInfo.ops);
-      processSingleNestedIf(
-          [](Directive nestedDir) { return topDistributeSet.test(nestedDir); });
+      processSingleNestedIf([](Directive nestedDir) {
+        return topDistributeSet.test(nestedDir) || topLoopSet.test(nestedDir);
+      });
       break;
 
     case OMPD_teams_distribute:
@@ -571,6 +572,12 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
       cp.processNumTeams(stmtCtx, hostInfo.ops);
       break;
 
+    case OMPD_loop:
+    case OMPD_teams_loop:
+    case OMPD_target_teams_loop:
+      cp.processCollapse(loc, eval, hostInfo.ops, hostInfo.iv);
+      break;
+
     // Standalone 'target' case.
     case OMPD_target:
       processSingleNestedIf(

diff --git a/flang/test/Lower/OpenMP/generic-loop-rewriting.f90 b/flang/test/Lower/OpenMP/generic-loop-rewriting.f90
@@ -11,7 +11,13 @@ subroutine target_teams_loop
 end subroutine target_teams_loop
 
 !CHECK-LABEL: func.func @_QPtarget_teams_loop
-!CHECK:         omp.target map_entries(
+!CHECK:         omp.target
+!CHECK-SAME:    host_eval(
+!CHECK-SAME:      %c0{{[^[:space:]]+}} -> %[[LB:[^[:space:]]+]],
+!CHECK-SAME:      %c10{{[^[:space:]]+}} -> %[[UB:[^[:space:]]+]],
+!CHECK-SAME:      %c1{{[^[:space:]]+}} -> %[[STEP:[^[:space:]]+]]
+!CHECK-SAME:      : i32, i32, i32)
+!CHECK-SAME:    map_entries(
 !CHECK-SAME:      %{{.*}} -> %[[I_ARG:[^[:space:]]+]],
 !CHECK-SAME:      %{{.*}} -> %[[X_ARG:[^[:space:]]+]] : {{.*}}) {
 
@@ -20,10 +26,6 @@ end subroutine target_teams_loop
 
 !CHECK:           omp.teams {
 
-!CHECK:             %[[LB:.*]] = arith.constant 0 : i32
-!CHECK:             %[[UB:.*]] = arith.constant 10 : i32
-!CHECK:             %[[STEP:.*]] = arith.constant 1 : i32
-
 !CHECK:             omp.parallel private(@{{.*}} %[[I_DECL]]#0 
 !CHECK-SAME:          -> %[[I_PRIV_ARG:[^[:space:]]+]] : !fir.ref<i32>) {
 !CHECK:               omp.distribute {

diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1909,7 +1909,7 @@ llvm::omp::OMPTgtExecModeFlags TargetOp::getKernelExecFlags() {
 
   // Detect Generic-SPMD: target-teams-distribute[-simd].
   if (numWrappers == 1) {
-    if (!isa<DistributeOp>(innermostWrapper))
+    if (!isa<DistributeOp>(innermostWrapper) && !isa<LoopOp>(innermostWrapper))
       return OMP_TGT_EXEC_MODE_GENERIC;
 
     Operation *teamsOp = (*innermostWrapper)->getParentOp();