Line data    Source code

       1            0 : // Distributed under the MIT License.
       2              : // See LICENSE.txt for details.
       3              : 
       4              : #pragma once
       5              : 
       6              : #include <cstddef>
       7              : #include <type_traits>
       8              : 
       9              : #include "DataStructures/DataVector.hpp"
      10              : #include "DataStructures/Tensor/Tensor.hpp"
      11              : #include "Utilities/Kokkos/KokkosCore.hpp"
      12              : 
      13              : #ifdef SPECTRE_KOKKOS
      14              : /// Copy a `Tensor<DataVector>` to device memory.
      15              : template <typename HostVectorType, typename... Properties>
      16            1 : auto copy_to_device(const Tensor<HostVectorType, Properties...>& tensor_host) {
      17              :   using ValueType = typename HostVectorType::value_type;
      18              :   using DeviceVectorType = Kokkos::View<ValueType*>;
      19              :   Tensor<DeviceVectorType, Properties...> tensor_device{
      20              :       "Tensor", tensor_host.begin()->size()};
      21              :   // View that wraps the host data (DataVector) in a Kokkos::View without owning
      22              :   // it so we can deep-copy the data
      23              :   using HostUnmanaged = Kokkos::View<const ValueType*, Kokkos::HostSpace,
      24              :                                      Kokkos::MemoryUnmanaged>;
      25              :   for (size_t i = 0; i < tensor_host.size(); ++i) {
      26              :     HostUnmanaged unmanaged_host_view(tensor_host[i].data(),
      27              :                                       tensor_host[i].size());
      28              :     // Copy to device
      29              :     Kokkos::deep_copy(tensor_device[i], unmanaged_host_view);
      30              :   }
      31              :   return tensor_device;
      32              : }
      33              : 
      34              : /// Copy a `Tensor<Kokkos::View>` to host memory.
      35              : template <typename DeviceVectorType, typename... Properties>
      36            1 : auto copy_to_host(
      37              :     const Tensor<DeviceVectorType, Properties...>& tensor_device) {
      38              :   using ValueType = typename DeviceVectorType::value_type;
      39              :   static_assert(
      40              :       std::is_same_v<ValueType, double>,
      41              :       "Only Tensor<Kokkos::View<double*>> is currently supported. Other "
      42              :       "datatypes can be added by mapping to their host representation (e.g. "
      43              :       "complex -> ComplexDataVector).");
      44              :   using HostVectorType = DataVector;
      45              :   Tensor<HostVectorType, Properties...> tensor_host{
      46              :       tensor_device.begin()->size()};
      47              :   // View that wraps the host data (DataVector) in a Kokkos::View without owning
      48              :   // it so we can deep-copy the data
      49              :   using HostUnmanaged =
      50              :       Kokkos::View<ValueType*, Kokkos::HostSpace, Kokkos::MemoryUnmanaged>;
      51              :   for (size_t i = 0; i < tensor_device.size(); ++i) {
      52              :     HostUnmanaged unmanaged_host_view(tensor_host[i].data(),
      53              :                                       tensor_host[i].size());
      54              :     // Copy to host
      55              :     Kokkos::deep_copy(unmanaged_host_view, tensor_device[i]);
      56              :   }
      57              :   return tensor_host;
      58              : }
      59              : #endif  // SPECTRE_KOKKOS