diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp index c89f577870eea7..0f7a7d1f52fdea 100644 --- a/flang/lib/Lower/OpenMP/OpenMP.cpp +++ b/flang/lib/Lower/OpenMP/OpenMP.cpp @@ -567,8 +567,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: @@ -581,6 +582,17 @@ static void processHostEvalClauses(lower::AbstractConverter &converter, cp.processNumTeams(stmtCtx, hostInfo.ops); break; + + case OMPD_teams_loop: + cp.processThreadLimit(stmtCtx, hostInfo.ops); + [[fallthrough]]; + case OMPD_target_teams_loop: + cp.processNumTeams(stmtCtx, hostInfo.ops); + [[fallthrough]]; + case OMPD_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 index fa26425356dd90..2c5ea367b7af90 100644 --- 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) { !CHECK: omp.distribute { diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp index c0b3abb87dfb3f..d879282a5999c8 100644 --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp @@ -1908,8 +1908,9 @@ llvm::omp::OMPTgtExecModeFlags TargetOp::getKernelExecFlags() { long numWrappers = std::distance(innermostWrapper, wrappers.end()); // Detect Generic-SPMD: target-teams-distribute[-simd]. + // Detect SPMD: target-teams-loop. if (numWrappers == 1) { - if (!isa(innermostWrapper)) + if (!isa(innermostWrapper)) return OMP_TGT_EXEC_MODE_GENERIC; Operation *teamsOp = (*innermostWrapper)->getParentOp(); @@ -1917,7 +1918,9 @@ llvm::omp::OMPTgtExecModeFlags TargetOp::getKernelExecFlags() { return OMP_TGT_EXEC_MODE_GENERIC; if (teamsOp->getParentOp() == *this) - return OMP_TGT_EXEC_MODE_GENERIC_SPMD; + return isa(innermostWrapper) + ? OMP_TGT_EXEC_MODE_GENERIC_SPMD + : OMP_TGT_EXEC_MODE_SPMD; } // Detect SPMD: target-teams-distribute-parallel-wsloop[-simd]. diff --git a/offload/test/offloading/fortran/target_teams_loop.f90 b/offload/test/offloading/fortran/target_teams_loop.f90 new file mode 100644 index 00000000000000..33670c79b5e32b --- /dev/null +++ b/offload/test/offloading/fortran/target_teams_loop.f90 @@ -0,0 +1,18 @@ +! REQUIRES: flang, amdgpu + +! RUN: %libomptarget-compile-fortran-generic +! RUN: env LIBOMPTARGET_INFO=16 %libomptarget-run-generic 2>&1 | %fcheck-generic +program target_teams_loop + implicit none + integer :: x(10), i + + !$omp target teams loop + do i = 1, 10 + x(i) = i * 2 + end do + + print *, x +end program target_teams_loop + +! CHECK: "PluginInterface" device {{[0-9]+}} info: Launching kernel {{.*}} +! CHECK: 2 4 6 8 10 12 14 16 18 20