ad threadvector

masterleinad · Oct 12, 2023 · 02e6bdc · 02e6bdc
1 parent 578bc7f
commit 02e6bdc
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/core/src/OpenMPTarget/Kokkos_OpenMPTarget_ParallelScan_Team.hpp b/core/src/OpenMPTarget/Kokkos_OpenMPTarget_ParallelScan_Team.hpp
@@ -139,6 +139,30 @@ KOKKOS_INLINE_FUNCTION void parallel_scan(
   }
 }
 
+template <typename iType, class FunctorType, class ValueType>
+KOKKOS_INLINE_FUNCTION void parallel_scan(
+    const Impl::ThreadVectorRangeBoundariesStruct<
+        iType, Impl::OpenMPTargetExecTeamMember>& loop_boundaries,
+    const FunctorType& lambda, ValueType& return_val) {
+  using Analysis = Impl::FunctorAnalysis<Impl::FunctorPatternInterface::SCAN,
+                                         TeamPolicy<Experimental::OpenMPTarget>,
+                                         FunctorType, void>;
+  using analysis_value_type = typename Analysis::value_type;
+  static_assert(std::is_same_v<analysis_value_type, ValueType>,
+                "Non-matching value types of functor and return type");
+
+  ValueType scan_val = {};
+
+#ifdef KOKKOS_ENABLE_PRAGMA_IVDEP
+#pragma ivdep
+#endif
+  for (iType i = loop_boundaries.start; i < loop_boundaries.end; ++i) {
+    lambda(i, scan_val, true);
+  }
+
+  return_val = scan_val;
+}
+
 }  // namespace Kokkos
 
 #ifdef KOKKOS_IMPL_TEAM_SCAN_WORKAROUND

diff --git a/core/unit_test/TestTeamVector.hpp b/core/unit_test/TestTeamVector.hpp
@@ -607,7 +607,7 @@ struct functor_vec_scan {
 
 // Temporary: This condition will progressively be reduced when parallel_scan
 // with return value will be implemented for more backends.
-#if !defined(KOKKOS_ENABLE_OPENACC) && !defined(KOKKOS_ENABLE_OPENMPTARGET)
+#if !defined(KOKKOS_ENABLE_OPENACC)
 template <typename Scalar, class ExecutionSpace>
 struct functor_vec_scan_ret_val {
   using policy_type     = Kokkos::TeamPolicy<ExecutionSpace>;
@@ -734,7 +734,7 @@ bool test_scalar(int nteams, int team_size, int test) {
   } else if (test == 12) {
 // Temporary: This condition will progressively be reduced when parallel_scan
 // with return value will be implemented for more backends.
-#if !defined(KOKKOS_ENABLE_OPENACC) && !defined(KOKKOS_ENABLE_OPENMPTARGET)
+#if !defined(KOKKOS_ENABLE_OPENACC)
     Kokkos::parallel_for(
         Kokkos::TeamPolicy<ExecutionSpace>(nteams, team_size, 8),
         functor_vec_scan_ret_val<Scalar, ExecutionSpace>(d_flag, team_size));