OpenMPTarget init-join fix (kokkos#6444)

* OpenMPTarget: Add verification in init join reduction. * OpenMPTarget: set MAX_ACTIVE_THREADS in initialize. * OpenMPTarget: fix error message. * OpenMPTarget: Fix error message. --------- Co-authored-by: Rahulkumar Gayatri <rgayatri@lbl.gov>
tcclevenger · Sep 19, 2023 · 7e27496 · 7e27496
1 parent 56cc35b
commit 7e27496
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 10 deletions.
diff --git a/core/src/OpenMPTarget/Kokkos_OpenMPTarget_Exec.cpp b/core/src/OpenMPTarget/Kokkos_OpenMPTarget_Exec.cpp
@@ -67,8 +67,6 @@ void OpenMPTargetExec::verify_initialized(const char* const label) {
     msg.append(" ERROR: not initialized");
     Kokkos::Impl::throw_runtime_exception(msg);
   }
-  OpenMPTargetExec::MAX_ACTIVE_THREADS =
-      Kokkos::Experimental::OpenMPTarget().concurrency();
 }
 
 void* OpenMPTargetExec::m_scratch_ptr         = nullptr;

diff --git a/core/src/OpenMPTarget/Kokkos_OpenMPTarget_Instance.cpp b/core/src/OpenMPTarget/Kokkos_OpenMPTarget_Instance.cpp
@@ -112,9 +112,12 @@ void OpenMPTargetInternal::impl_finalize() {
     Kokkos::kokkos_free<Kokkos::Experimental::OpenMPTargetSpace>(
         space.m_uniquetoken_ptr);
 }
+
 void OpenMPTargetInternal::impl_initialize() {
   m_is_initialized = true;
 
+  Kokkos::Impl::OpenMPTargetExec::MAX_ACTIVE_THREADS = concurrency();
+
   // FIXME_OPENMPTARGET:  Only fix the number of teams for NVIDIA architectures
   // from Pascal and upwards.
   // FIXME_OPENMPTARGTE: Cray compiler did not yet implement omp_set_num_teams.

diff --git a/core/src/OpenMPTarget/Kokkos_OpenMPTarget_Parallel_Common.hpp b/core/src/OpenMPTarget/Kokkos_OpenMPTarget_Parallel_Common.hpp
@@ -84,9 +84,11 @@ struct ParallelReduceSpecialize<FunctorType, Kokkos::RangePolicy<PolicyArgs...>,
   static void execute_reducer(const FunctorType& f, const PolicyType& p,
                               PointerType result_ptr, bool ptr_on_device) {
     OpenMPTargetExec::verify_is_process(
-        "Kokkos::Experimental::OpenMPTarget parallel_for");
+        "Kokkos::Experimental::OpenMPTarget RangePolicy "
+        "parallel_reduce:reducer");
     OpenMPTargetExec::verify_initialized(
-        "Kokkos::Experimental::OpenMPTarget parallel_for");
+        "Kokkos::Experimental::OpenMPTarget RangePolicy "
+        "parallel_reduce:reducer");
     const auto begin = p.begin();
     const auto end   = p.end();
 
@@ -125,9 +127,11 @@ struct ParallelReduceSpecialize<FunctorType, Kokkos::RangePolicy<PolicyArgs...>,
   static void execute_array(const FunctorType& f, const PolicyType& p,
                             PointerType result_ptr, bool ptr_on_device) {
     OpenMPTargetExec::verify_is_process(
-        "Kokkos::Experimental::OpenMPTarget parallel_for");
+        "Kokkos::Experimental::OpenMPTarget RangePolicy "
+        "parallel_reduce:array_reduction");
     OpenMPTargetExec::verify_initialized(
-        "Kokkos::Experimental::OpenMPTarget parallel_for");
+        "Kokkos::Experimental::OpenMPTarget RangePolicy "
+        "parallel_reduce:array_reduction");
     const auto begin = p.begin();
     const auto end   = p.end();
 
@@ -197,6 +201,12 @@ struct ParallelReduceSpecialize<FunctorType, Kokkos::RangePolicy<PolicyArgs...>,
 
   static void execute_init_join(const FunctorType& f, const PolicyType& p,
                                 PointerType ptr, const bool ptr_on_device) {
+    OpenMPTargetExec::verify_is_process(
+        "Kokkos::Experimental::OpenMPTarget RangePolicy "
+        "parallel_reduce:init_join");
+    OpenMPTargetExec::verify_initialized(
+        "Kokkos::Experimental::OpenMPTarget RangePolicy "
+        "parallel_reduce:init_join");
     const auto begin = p.begin();
     const auto end   = p.end();
 
@@ -347,9 +357,11 @@ struct ParallelReduceSpecialize<FunctorType, TeamPolicyInternal<PolicyArgs...>,
   static void execute_reducer(const FunctorType& f, const PolicyType& p,
                               PointerType result_ptr, bool ptr_on_device) {
     OpenMPTargetExec::verify_is_process(
-        "Kokkos::Experimental::OpenMPTarget parallel_for");
+        "Kokkos::Experimental::OpenMPTarget TeamPolicy "
+        "parallel_reduce:reducer");
     OpenMPTargetExec::verify_initialized(
-        "Kokkos::Experimental::OpenMPTarget parallel_for");
+        "Kokkos::Experimental::OpenMPTarget TeamPolicy "
+        "parallel_reduce:reducer");
 
     const int league_size   = p.league_size();
     const int team_size     = p.team_size();
@@ -435,9 +447,11 @@ struct ParallelReduceSpecialize<FunctorType, TeamPolicyInternal<PolicyArgs...>,
   static void execute_array(const FunctorType& f, const PolicyType& p,
                             PointerType result_ptr, bool ptr_on_device) {
     OpenMPTargetExec::verify_is_process(
-        "Kokkos::Experimental::OpenMPTarget parallel_for");
+        "Kokkos::Experimental::OpenMPTarget TeamPolicy "
+        "parallel_reduce:array_reduction");
     OpenMPTargetExec::verify_initialized(
-        "Kokkos::Experimental::OpenMPTarget parallel_for");
+        "Kokkos::Experimental::OpenMPTarget TeamPolicy "
+        "parallel_reduce:array_reduction");
 
     const int league_size   = p.league_size();
     const int team_size     = p.team_size();
@@ -560,6 +574,12 @@ struct ParallelReduceSpecialize<FunctorType, TeamPolicyInternal<PolicyArgs...>,
   // RangePolicy. Need a new implementation.
   static void execute_init_join(const FunctorType& f, const PolicyType& p,
                                 PointerType ptr, const bool ptr_on_device) {
+    OpenMPTargetExec::verify_is_process(
+        "Kokkos::Experimental::OpenMPTarget TeamPolicy "
+        "parallel_reduce:init_join ");
+    OpenMPTargetExec::verify_initialized(
+        "Kokkos::Experimental::OpenMPTarget TeamPolicy "
+        "parallel_reduce:init_join");
     using FunctorAnalysis =
         Impl::FunctorAnalysis<Impl::FunctorPatternInterface::REDUCE, PolicyType,
                               FunctorType, ValueType>;