From fd1d45ae826d3e8c1a0ae8437958f451092684be Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Wed, 30 Nov 2022 16:36:29 +0000 Subject: [PATCH] Update memory_space --- include/deal.II/base/memory_space.h | 19 ++- include/deal.II/base/memory_space_data.h | 193 ++++++++++------------- 2 files changed, 99 insertions(+), 113 deletions(-) diff --git a/include/deal.II/base/memory_space.h b/include/deal.II/base/memory_space.h index 8c1f10ac29..34db769f22 100644 --- a/include/deal.II/base/memory_space.h +++ b/include/deal.II/base/memory_space.h @@ -19,6 +19,8 @@ #include +#include + DEAL_II_NAMESPACE_OPEN /** @@ -29,15 +31,24 @@ namespace MemorySpace * Structure describing Host memory space. */ struct Host - {}; - + { + using kokkos_space = ::Kokkos::HostSpace; + }; + /** + * Structure describing Device memory space. + */ + struct Device + { + using kokkos_space = ::Kokkos::DefaultExecutionSpace::memory_space; + }; +#ifdef DEAL_II_WITH_CUDA /** * Structure describing CUDA memory space. */ - struct CUDA - {}; + using CUDA = Device; +#endif } // namespace MemorySpace diff --git a/include/deal.II/base/memory_space_data.h b/include/deal.II/base/memory_space_data.h index e2feb375db..da0b809c8d 100644 --- a/include/deal.II/base/memory_space_data.h +++ b/include/deal.II/base/memory_space_data.h @@ -21,6 +21,9 @@ #include #include +#include + +#include #include #include @@ -32,54 +35,55 @@ DEAL_II_NAMESPACE_OPEN namespace MemorySpace { /** - * Data structure + * Structure which stores data on the host or the device depending on the + * template parameter @p MemorySpace. The data is copied into the structure + * which then owns the data and will release the memory when the destructor is + * called. */ - template + template struct MemorySpaceData { - MemorySpaceData() - { - static_assert(std::is_same::value || - std::is_same::value, - "MemorySpace should be Host or CUDA"); - } + MemorySpaceData(); /** - * Copy the active data (values for Host and values_dev for CUDA) to @p begin. + * Copy the active data (values for Host and values_dev for Device) to @p begin. * If the data is on the device it is moved to the host. */ void - copy_to(Number *begin, std::size_t n_elements) - { - (void)begin; - (void)n_elements; - } + copy_to(T *begin, const std::size_t n_elements); /** * Copy the data in @p begin to the active data of the structure (values for - * Host and values_dev for CUDA). The pointer @p begin must be on the host. + * Host and values_dev for Device). The pointer @p begin must be on the host. */ void - copy_from(Number *begin, std::size_t n_elements) - { - (void)begin; - (void)n_elements; - } + copy_from(const T *begin, const std::size_t n_elements); + + /** + * Kokkos View to the data on the host + */ + Kokkos::View values; /** - * Pointer to data on the host. + * Kokkos View to the data on the device */ - std::unique_ptr> values; + Kokkos::View values_dev; /** - * Pointer to data on the device. + * Pointer to data on the host. The pointer points to the same data as + * values when using shared memory. Otherwise it is not set. */ - std::unique_ptr values_dev; + // The pointer is shared 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 in non-owning. When shared memory with MPI is not used, + // the pointer is not used. + std::shared_ptr values_sm_ptr; /** * Pointers to the data of the processes sharing the same memory. */ - std::vector> values_sm; + std::vector> values_sm; }; @@ -87,110 +91,81 @@ namespace MemorySpace /** * Swap function similar to std::swap. */ - template + template inline void - swap(MemorySpaceData &, - MemorySpaceData &) - { - static_assert(std::is_same::value || - std::is_same::value, - "MemorySpace should be Host or CUDA"); - } - -#ifndef DOXYGEN + swap(MemorySpaceData &u, MemorySpaceData &v); - template - struct MemorySpaceData - { - MemorySpaceData() - : values(nullptr, &std::free) - {} - - void - copy_to(Number *begin, std::size_t n_elements) - { - std::copy(values.get(), values.get() + n_elements, begin); - } - - void - copy_from(Number *begin, std::size_t n_elements) - { - std::copy(begin, begin + n_elements, values.get()); - } - - std::unique_ptr> values; - // This is not used but it allows to simplify the code until we start using - // CUDA-aware MPI. - std::unique_ptr values_dev; +#ifndef DOXYGEN - std::vector> values_sm; - }; + template + MemorySpaceData::MemorySpaceData() + : values((dealii::Impl::ensure_kokkos_initialized(), + Kokkos::View("host data", 0))) + , values_dev(Kokkos::View( + "memoryspace data", + 0)) + {} - template - inline void - swap(MemorySpaceData &u, MemorySpaceData &v) + template + void + MemorySpaceData::copy_to(T * begin, + const std::size_t n_elements) { - std::swap(u.values, v.values); + Assert(n_elements <= values.extent(0), + ExcMessage("n_elements greater than the size of values.")); + using ExecutionSpace = typename MemorySpace::kokkos_space::execution_space; + Kokkos:: + View> + 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(); } -# ifdef DEAL_II_COMPILER_CUDA_AWARE - - template - struct MemorySpaceData + template + void + MemorySpaceData::copy_from(const T * begin, + const std::size_t n_elements) { - MemorySpaceData() - : values(nullptr, &std::free) - , values_dev(nullptr, Utilities::CUDA::delete_device_data) - {} - - void - copy_to(Number *begin, std::size_t n_elements) - { - const cudaError_t cuda_error_code = - cudaMemcpy(begin, - values_dev.get(), - n_elements * sizeof(Number), - cudaMemcpyDeviceToHost); - AssertCuda(cuda_error_code); - } - - void - copy_from(Number *begin, std::size_t n_elements) - { - const cudaError_t cuda_error_code = - cudaMemcpy(values_dev.get(), - begin, - n_elements * sizeof(Number), - cudaMemcpyHostToDevice); - AssertCuda(cuda_error_code); - } - - std::unique_ptr> values; - std::unique_ptr values_dev; - - /** - * This is currently not used. - */ - std::vector> values_sm; - }; + Assert(n_elements <= values.extent(0), + ExcMessage("n_elements greater than the size of values.")); + using ExecutionSpace = typename MemorySpace::kokkos_space::execution_space; + Kokkos::View> + 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 + /** + * Swap function similar to std::swap. + */ + template inline void - swap(MemorySpaceData &u, MemorySpaceData &v) + swap(MemorySpaceData &u, MemorySpaceData &v) { - std::swap(u.values, v.values); - std::swap(u.values_dev, v.values_dev); + auto u_copy = Kokkos::create_mirror(Kokkos::WithoutInitializing, u); + typename MemorySpace::execution_space exec_space; + // The first two calls to Kokkos::deep_copy are asynchronous. The last call + // will wait for the three deep_copy to be done before returning. + Kokkos::deep_copy(exec_space, u_copy, u); + Kokkos::deep_copy(exec_space, u, v); + Kokkos::deep_copy(v, u_copy); } -# endif - #endif } // namespace MemorySpace -- 2.39.5