memory_space_data.h

Go to the documentation of this file.

// ---------------------------------------------------------------------
//
// Copyright (C) 2020 - 2023 by the deal.II authors
//
// This file is part of the deal.II library.
//
// The deal.II library is free software; you can use it, redistribute
// it, and/or modify it under the terms of the GNU Lesser General
// Public License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
// The full text of the license can be found in the file LICENSE.md at
// the top level directory of deal.II.
//
// ---------------------------------------------------------------------
 
 
#ifndef dealii_memory_space_data_h
#define dealii_memory_space_data_h
 
#include <deal.II/base/config.h>
 
#include <deal.II/base/cuda.h>
#include <deal.II/base/exceptions.h>
#include <deal.II/base/kokkos.h>
 
#include <Kokkos_Core.hpp>
 
#include <functional>
#include <memory>
 
DEAL_II_NAMESPACE_OPEN
 
namespace MemorySpace
{
  template <typename T, typename MemorySpace>
  struct MemorySpaceData
  {
    MemorySpaceData();
 
    void
    copy_to(T *begin, const std::size_t n_elements);
 
    void
    copy_from(const T *begin, const std::size_t n_elements);
 
    // FIXME Should we move this somewhere else?
    Kokkos::View<T *, Kokkos::HostSpace> values_host_buffer;
 
    Kokkos::View<T *, typename MemorySpace::kokkos_space> values;
 
    // This a shared pointer so that MemorySpaceData can be copied and
    // MemorySpaceData::values can be used in Kokkos::parallel_for. This
    // pointer owns the data when using shared memory with MPI. In this case,
    // the Kokkos::View @p values is non-owning. When shared memory with MPI is
    // not used, the @p values_sm_ptr is unused.
    std::shared_ptr<T> values_sm_ptr;
 
    std::vector<ArrayView<const T>> values_sm;
  };
 
 
 
  template <typename T, typename MemorySpace>
  inline void
  swap(MemorySpaceData<T, MemorySpace> &u, MemorySpaceData<T, MemorySpace> &v);
 
 
#ifndef DOXYGEN
 
  template <typename T, typename MemorySpace>
  MemorySpaceData<T, MemorySpace>::MemorySpaceData()
    : values_host_buffer(
        (::internal::ensure_kokkos_initialized(),
         Kokkos::View<T *, Kokkos::HostSpace>("host buffer", 0)))
    , values(Kokkos::View<T *, typename MemorySpace::kokkos_space>(
        "memoryspace data",
        0))
  {}
 
 
 
  template <typename T, typename MemorySpace>
  void
  MemorySpaceData<T, MemorySpace>::copy_to(T *               begin,
                                           const std::size_t n_elements)
  {
    Assert(n_elements <= values.extent(0),
           ExcMessage(
             "n_elements is greater than the size of MemorySpaceData."));
    using ExecutionSpace = typename MemorySpace::kokkos_space::execution_space;
    Kokkos::
      View<T *, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>>
        begin_view(begin, n_elements);
    Kokkos::deep_copy(
      ExecutionSpace{},
      begin_view,
      Kokkos::subview(values, Kokkos::make_pair(std::size_t(0), n_elements)));
    ExecutionSpace{}.fence();
  }
 
 
 
  template <typename T, typename MemorySpace>
  void
  MemorySpaceData<T, MemorySpace>::copy_from(const T *         begin,
                                             const std::size_t n_elements)
  {
    Assert(n_elements <= values.extent(0),
           ExcMessage(
             "n_elements is greater than the size of MemorySpaceData."));
    using ExecutionSpace = typename MemorySpace::kokkos_space::execution_space;
    Kokkos::View<const T *,
                 Kokkos::HostSpace,
                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>
      begin_view(begin, n_elements);
    Kokkos::deep_copy(
      ExecutionSpace{},
      Kokkos::subview(values, Kokkos::make_pair(std::size_t(0), n_elements)),
      begin_view);
    ExecutionSpace{}.fence();
  }
 
 
 
  template <typename T, typename MemorySpace>
  inline void
  swap(MemorySpaceData<T, MemorySpace> &u, MemorySpaceData<T, MemorySpace> &v)
  {
    std::swap(u.values_host_buffer, v.values_host_buffer);
    std::swap(u.values, v.values);
    std::swap(u.values_sm_ptr, v.values_sm_ptr);
  }
 
#endif
 
} // namespace MemorySpace
 
DEAL_II_NAMESPACE_CLOSE
 
#endif