Skip to content

Commit

Permalink
#2281: Fix runtime Kokkos when we run with both device/host memory sp…
Browse files Browse the repository at this point in the history
…aces
  • Loading branch information
JacobDomagala committed Sep 20, 2024
1 parent ea24fc1 commit 897f570
Show file tree
Hide file tree
Showing 6 changed files with 151 additions and 19 deletions.
2 changes: 1 addition & 1 deletion src/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
42 changes: 42 additions & 0 deletions src/vt/collective/reduce/allreduce/allreduce_holder.cc
Original file line number Diff line number Diff line change
@@ -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 [email protected]
//
// *****************************************************************************
//@HEADER
*/
#include "allreduce_holder.h"
#include "vt/objgroup/manager.h"

Expand Down
22 changes: 16 additions & 6 deletions src/vt/collective/reduce/allreduce/helpers.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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>
Expand Down Expand Up @@ -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>
Expand All @@ -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));
}
);
Expand All @@ -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));
}
);
Expand Down
25 changes: 16 additions & 9 deletions src/vt/collective/reduce/operators/functors/plus_op.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Expand All @@ -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
Expand Down
73 changes: 73 additions & 0 deletions src/vt/utils/kokkos/exec_space.h
Original file line number Diff line number Diff line change
@@ -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 [email protected]
//
// *****************************************************************************
//@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*/
6 changes: 3 additions & 3 deletions tests/perf/allreduce.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down

0 comments on commit 897f570

Please sign in to comment.