Skip to content

Commit

Permalink
feat(kokkos): use Kokkos unmanaged views to avoid extra copy
Browse files Browse the repository at this point in the history
  • Loading branch information
pranav-sivaraman committed Jan 25, 2024
1 parent bea0523 commit d5f58d5
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 8 deletions.
14 changes: 7 additions & 7 deletions src/kokkos/fasten.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -172,10 +172,9 @@ template <size_t PPWI> class IMPL_CLS final : public Bude<PPWI> {
}

template <typename T> static Kokkos::View<T *> mkView(const std::string &name, const std::vector<T> &xs) {
Kokkos::View<T *> view(Kokkos::ViewAllocateWithoutInitializing(name), xs.size());
auto mirror = Kokkos::create_mirror_view(view);
for (size_t i = 0; i < xs.size(); i++)
mirror[i] = xs[i];
Kokkos::View<const T *, Kokkos::LayoutLeft, Kokkos::HostSpace,
Kokkos::MemoryTraits<Kokkos::Unmanaged>> mirror (std::data(xs), std::size(xs));
Kokkos::View<T *> view (name, std::size(xs));
Kokkos::deep_copy(view, mirror);
return view;
}
Expand Down Expand Up @@ -230,12 +229,13 @@ template <size_t PPWI> class IMPL_CLS final : public Bude<PPWI> {

auto result_mirror = Kokkos::create_mirror_view(results);
Kokkos::deep_copy(result_mirror, results);
for (size_t i = 0; i < results.size(); i++) {
sample.energies[i] = result_mirror[i];
}

auto deviceToHostEnd = now();
sample.deviceToHost = {deviceToHostStart, deviceToHostEnd};

for (size_t i = 0; i < results.size(); i++) {
sample.energies[i] = result_mirror[i];
}
}

if (!Kokkos::is_finalized()) {
Expand Down
1 change: 0 additions & 1 deletion src/kokkos/model.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@ macro(setup)
register_link_library(Kokkos::kokkos)
elseif (KOKKOS_IN_PACKAGE)
message(STATUS "Build using packaged Kokkos at `${KOKKOS_IN_PACKAGE}`")
set (Kokkos_DIR "${KOKKOS_IN_PACKAGE}/lib64/cmake/Kokkos")
find_package(Kokkos REQUIRED)
register_link_library(Kokkos::kokkos)
else ()
Expand Down

0 comments on commit d5f58d5

Please sign in to comment.