#2281: Fix runtime Kokkos when we run with both device/host memory sp…

…aces

src/CMakeLists.txt

@@ -99,7 +99,7 @@ set(
       serialization/messaging serialization/traits serialization/auto_dispatch
       serialization/sizing
     utils/demangle utils/container utils/bits utils/mutex utils/file_spec
-    utils/hash utils/atomic utils/static_checks utils/string
+    utils/hash utils/atomic utils/static_checks utils/string utils/kokkos
     utils/memory utils/mpi_limits utils/compress utils/json utils/strong
     registry/auto
       registry/auto/functor registry/auto/map registry/auto/collection

src/vt/collective/reduce/allreduce/allreduce_holder.cc

@@ -1,3 +1,45 @@
+/*
+//@HEADER
+// *****************************************************************************
+//
+//                             allreduce_holder.cc
+//                       DARMA/vt => Virtual Transport
+//
+// Copyright 2019-2021 National Technology & Engineering Solutions of Sandia, LLC
+// (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S.
+// Government retains certain rights in this software.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// * Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+//
+// * Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// * Neither the name of the copyright holder nor the names of its
+//   contributors may be used to endorse or promote products derived from this
+//   software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+// POSSIBILITY OF SUCH DAMAGE.
+//
+// Questions? Contact darma@sandia.gov
+//
+// *****************************************************************************
+//@HEADER
+*/
 #include "allreduce_holder.h"
 #include "vt/objgroup/manager.h"
 

src/vt/collective/reduce/allreduce/helpers.h

@@ -47,6 +47,8 @@
 #include "data_handler.h"
 #include "rabenseifner_msg.h"
 #include "vt/messaging/message/shared_message.h"
+#include "vt/utils/kokkos/exec_space.h"
+
 #include <vector>
 
 namespace vt::collective::reduce::allreduce {
@@ -61,6 +63,7 @@ template <typename Scalar>
 struct ShouldUseView<Scalar, Kokkos::View<Scalar*, Kokkos::HostSpace>> {
   static constexpr bool Value = true;
 };
+
 #endif // MAGISTRATE_KOKKOS_ENABLED
 
 template <typename Scalar, typename DataT>
@@ -116,9 +119,10 @@ struct DataHelper {
 
 #if MAGISTRATE_KOKKOS_ENABLED
 
-template <typename Scalar>
-struct DataHelper<Scalar, Kokkos::View<Scalar*, Kokkos::HostSpace>> {
-  using DataT = Kokkos::View<Scalar*, Kokkos::HostSpace>;
+template <typename Scalar, typename MemorySpace>
+struct DataHelper<Scalar, Kokkos::View<Scalar*, MemorySpace>> {
+  using DataT = Kokkos::View<Scalar*, MemorySpace>;
+  using ExecSpace = typename utils::kokkos::AssociatedExecSpace<MemorySpace>::type;
   using DataType = DataHandler<DataT>;
 
   template <typename... Args>
@@ -136,17 +140,21 @@ struct DataHelper<Scalar, Kokkos::View<Scalar*, Kokkos::HostSpace>> {
 
   static void
   copy(DataT& dest, size_t start_idx, RabenseifnerMsg<Scalar, DataT>* msg) {
+
+    Kokkos::RangePolicy<ExecSpace> policy(0, msg->val_.extent(0));
     Kokkos::parallel_for(
-      "Rabenseifner::copy", msg->val_.extent(0),
+      "Rabenseifner::copy", policy,
       KOKKOS_LAMBDA(const int i) { dest(start_idx + i) = msg->val_(i); }
     );
   }
 
   template <template <typename Arg> class Op>
   static void reduceMsg(
     DataT& dest, size_t start_idx, RabenseifnerMsg<Scalar, DataT>* msg) {
+
+    Kokkos::RangePolicy<ExecSpace> policy(0, msg->val_.extent(0));
     Kokkos::parallel_for(
-      "Rabenseifner::reduce", msg->val_.extent(0), KOKKOS_LAMBDA(const int i) {
+      "Rabenseifner::reduce", policy, KOKKOS_LAMBDA(const int i) {
         Op<Scalar>()(dest(start_idx + i), msg->val_(i));
       }
     );
@@ -156,8 +164,10 @@ struct DataHelper<Scalar, Kokkos::View<Scalar*, Kokkos::HostSpace>> {
   static void reduce(
     DataT& dest, Args&&... val) {
     auto view_val = DataT{std::forward<Args>(val)...};
+
+    Kokkos::RangePolicy<ExecSpace> policy(0, view_val.extent(0));
     Kokkos::parallel_for(
-      "Rabenseifner::reduce", view_val.extent(0), KOKKOS_LAMBDA(const int i) {
+      "Rabenseifner::reduce", policy, KOKKOS_LAMBDA(const int i) {
         Op<Scalar>()(dest(i), view_val(i));
       }
     );

src/vt/collective/reduce/operators/functors/plus_op.h

@@ -47,9 +47,7 @@
 #include "vt/config.h"
 #include "vt/collective/reduce/operators/functors/tuple_op_helper.h"
 
-#if MAGISTRATE_KOKKOS_ENABLED
-#include <Kokkos_Core.hpp>
-#endif
+#include "vt/utils/kokkos/exec_space.h"
 
 namespace vt { namespace collective { namespace reduce { namespace operators {
 
@@ -71,14 +69,23 @@ struct PlusOp<std::tuple<Params...>> {
 
 #if MAGISTRATE_KOKKOS_ENABLED
 
-template <typename T>
-struct PlusOp<Kokkos::View<T*, Kokkos::HostSpace>> {
+template <typename T, typename MemorySpace>
+struct PlusOp<Kokkos::View<T*, MemorySpace>> {
   void operator()(
-    Kokkos::View<T*, Kokkos::HostSpace>& v1,
-    Kokkos::View<T*, Kokkos::HostSpace> const& v2) {
+    Kokkos::View<T*, MemorySpace>& v1,
+    Kokkos::View<T*, MemorySpace> const& v2) const {
+
+    using ExecSpace = typename utils::kokkos::AssociatedExecSpace<MemorySpace>::type;
+
+    Kokkos::RangePolicy<ExecSpace> policy(0, v1.extent(0));
+
     Kokkos::parallel_for(
-      "Initialize", v1.extent(0),
-      KOKKOS_LAMBDA(const int i) { v1(i) += v2(i); });
+      "PlusOp_Host",
+      policy,
+      KOKKOS_LAMBDA(const int i) {
+        v1(i) += v2(i);
+      }
+    );
   }
 };
 #endif // MAGISTRATE_KOKKOS_ENABLED

src/vt/utils/kokkos/exec_space.h

@@ -0,0 +1,73 @@
+/*
+//@HEADER
+// *****************************************************************************
+//
+//                                 exec_space.h
+//                       DARMA/vt => Virtual Transport
+//
+// Copyright 2019-2021 National Technology & Engineering Solutions of Sandia, LLC
+// (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S.
+// Government retains certain rights in this software.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// * Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+//
+// * Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// * Neither the name of the copyright holder nor the names of its
+//   contributors may be used to endorse or promote products derived from this
+//   software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+// POSSIBILITY OF SUCH DAMAGE.
+//
+// Questions? Contact darma@sandia.gov
+//
+// *****************************************************************************
+//@HEADER
+*/
+
+#if !defined INCLUDED_VT_UTILS_KOKKOS_EXEC_SPACE_H
+#define INCLUDED_VT_UTILS_KOKKOS_EXEC_SPACE_H
+
+#include "vt/config.h"
+
+#if MAGISTRATE_KOKKOS_ENABLED
+
+#include <Kokkos_Core.hpp>
+
+namespace vt::utils::kokkos {
+
+template <typename MemorySpace>
+struct AssociatedExecSpace;
+
+template <>
+struct AssociatedExecSpace<Kokkos::HostSpace> {
+  // TODO: check whether OpenMPi is available?
+  using type = Kokkos::Serial;
+};
+
+#ifdef KOKKOS_ENABLE_CUDA
+template <>
+struct AssociatedExecSpace<Kokkos::CudaSpace> {
+  using type = Kokkos::Cuda;
+};
+#endif // KOKKOS_ENABLE_CUDA
+}
+
+#endif // MAGISTRATE_KOKKOS_ENABLED
+#endif /*INCLUDED_VT_UTILS_KOKKOS_EXEC_SPACE_H*/

tests/perf/allreduce.cc

@@ -141,9 +141,9 @@ VT_PERF_TEST(MyTestKokkos, test_reduce_kokkos) {
 
   for (auto payload_size : payloadSizes) {
     view = Kokkos::View<float*, Kokkos::HostSpace>("view", payload_size);
-    Kokkos::parallel_for(
-      "InitView", view.extent(0),
-      KOKKOS_LAMBDA(const int i) { view(i) = static_cast<float>(my_node_); });
+    for (uint32_t i = 0; i < view.extent(0); ++i) {
+      view(i) = static_cast<float>(my_node_);
+    }
 
     theCollective()->barrier();
     auto* obj_ptr = grp_proxy[my_node_].get();