Fix reducer result check for Threads ParallelReduce

sandialabs · Mar 2, 2023 · 7b598eb · 7b598eb
1 parent 9a33347
commit 7b598eb
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 7 deletions.
diff --git a/core/src/Threads/Kokkos_Threads_Parallel_MDRange.hpp b/core/src/Threads/Kokkos_Threads_Parallel_MDRange.hpp
@@ -220,8 +220,10 @@ class ParallelReduce<CombinedFunctorReducerType,
                   "Kokkos::Threads reduce result must be a View");
 
     static_assert(
-        std::is_same<typename HostViewType::memory_space, HostSpace>::value,
-        "Kokkos::Threads reduce result must be a View in HostSpace");
+        Kokkos::Impl::MemorySpaceAccess<typename ViewType::memory_space,
+                                        Kokkos::HostSpace>::accessible,
+        "Kokkos::Threads reduce result must be a View accessible from "
+        "HostSpace");
   }
 
   template <typename Policy, typename Functor>

diff --git a/core/src/Threads/Kokkos_Threads_Parallel_Range.hpp b/core/src/Threads/Kokkos_Threads_Parallel_Range.hpp
@@ -242,18 +242,20 @@ class ParallelReduce<CombinedFunctorReducerType, Kokkos::RangePolicy<Traits...>,
     }
   }
 
-  template <class HostViewType>
+  template <class ViewType>
   ParallelReduce(const CombinedFunctorReducerType &arg_functor_reducer,
-                 const Policy &arg_policy, const HostViewType &arg_result_view)
+                 const Policy &arg_policy, const ViewType &arg_result_view)
       : m_functor_reducer(arg_functor_reducer),
         m_policy(arg_policy),
         m_result_ptr(arg_result_view.data()) {
-    static_assert(Kokkos::is_view<HostViewType>::value,
+    static_assert(Kokkos::is_view<ViewType>::value,
                   "Kokkos::Threads reduce result must be a View");
 
     static_assert(
-        std::is_same<typename HostViewType::memory_space, HostSpace>::value,
-        "Kokkos::Threads reduce result must be a View in HostSpace");
+        Kokkos::Impl::MemorySpaceAccess<typename ViewType::memory_space,
+                                        Kokkos::HostSpace>::accessible,
+        "Kokkos::Threads reduce result must be a View accessible from "
+        "HostSpace");
   }
 };
 

diff --git a/core/src/Threads/Kokkos_Threads_Parallel_Team.hpp b/core/src/Threads/Kokkos_Threads_Parallel_Team.hpp
@@ -212,6 +212,11 @@ class ParallelReduce<CombinedFunctorReducerType,
         m_shared(m_policy.scratch_size(0) + m_policy.scratch_size(1) +
                  FunctorTeamShmemSize<FunctorType>::value(
                      arg_functor_reducer.get_functor(), m_policy.team_size())) {
+    static_assert(
+        Kokkos::Impl::MemorySpaceAccess<typename ViewType::memory_space,
+                                        Kokkos::HostSpace>::accessible,
+        "Kokkos::Threads reduce result must be a View accessible from "
+        "HostSpace");
   }
 };