Line data    Source code

       1            0 : // Distributed under the MIT License.
       2              : // See LICENSE.txt for details.
       3              : 
       4              : #pragma once
       5              : 
       6              : #include <utility>
       7              : 
       8              : #include "DataStructures/DataBox/Tag.hpp"
       9              : #include "DataStructures/Tensor/Metafunctions.hpp"
      10              : #include "DataStructures/Tensor/Tensor.hpp"
      11              : #include "Utilities/Kokkos/KokkosCore.hpp"
      12              : 
      13              : #ifdef SPECTRE_KOKKOS
      14              : 
      15              : /// \cond
      16              : class DataVector;
      17              : /// \endcond
      18              : 
      19              : namespace Tags {
      20              : 
      21              : /*!
      22              :  * \brief Prefix tag representing the given tensor tag in a Kokkos memory space
      23              :  *
      24              :  * Given a `Tag` that holds a `Tensor`, this tag represents the same tensor in
      25              :  * the given Kokkos memory space, or in the default Kokkos memory space if
      26              :  * unspecified (see Kokkos documentation on memory spaces). Examples:
      27              :  * - Given a `Tag` that holds a `Tensor<DataVector>`, this tag holds a
      28              :  *   `Tensor<Kokkos::View<double*>>`, which is allocated in device memory by
      29              :  *   default. This is useful to mirror non-Kokkos data to device memory (see
      30              :  *   `copy_to_device`).
      31              :  * - Given a `Tag` that holds a `Tensor<Kokkos::View<double*>>` and specifying
      32              :  *   `Space = Kokkos::HostSpace`, this tag holds a host mirror of the device
      33              :  *   data. This is useful for creating a mirror view of the tensor data in host
      34              :  *   memory to move data around.
      35              :  */
      36              : template <typename Tag,
      37              :           typename Space = typename Kokkos::DefaultExecutionSpace::memory_space,
      38              :           typename VectorType = typename Tag::type::type>
      39            1 : struct MirrorView;
      40              : 
      41              : /// \cond
      42              : // Specialization for `DataVector` -> mirror to `Kokkos::View<double*>`
      43              : // (Omit `Space` in `Kokkos::View<double*>` because that would make it a
      44              : // different type. Plain `Kokkos::View<double*>` is what would be hand-written
      45              : // or returned by `copy_to_device`. We can add a specialization if custom memory
      46              : // spaces are needed.)
      47              : template <typename Tag>
      48              : struct MirrorView<Tag, typename Kokkos::DefaultExecutionSpace::memory_space,
      49              :                   DataVector> : db::PrefixTag,
      50              :                                 db::SimpleTag {
      51              :   using tag = Tag;
      52              :   using type =
      53              :       TensorMetafunctions::swap_type<Kokkos::View<double*>, typename Tag::type>;
      54              : };
      55              : 
      56              : template <typename Tag, typename Space, typename DataType,
      57              :           typename... Properties>
      58              : struct MirrorView<Tag, Space, Kokkos::View<DataType, Properties...>>
      59              :     : db::PrefixTag, db::SimpleTag {
      60              :   using tag = Tag;
      61              :   using type = TensorMetafunctions::swap_type<
      62              :       // Use public Kokkos API to get the mirror view type
      63              :       decltype(Kokkos::create_mirror_view(
      64              :           std::declval<Space>(),
      65              :           std::declval<Kokkos::View<DataType, Properties...>>())),
      66              :       typename Tag::type>;
      67              : };
      68              : /// \endcond
      69              : 
      70              : }  // namespace Tags
      71              : 
      72              : #endif  // SPECTRE_KOKKOS