Line data    Source code

       1              : //
       2              : // Class Archive
       3              : //   Class to (de-)serialize in MPI communication.
       4              : //
       5              : #include <cstring>
       6              : 
       7              : #include "Archive.h"
       8              : 
       9              : namespace ippl {
      10              :     namespace detail {
      11              : 
      12              :         template <class... Properties>
      13           26 :         Archive<Properties...>::Archive(size_type size)
      14           26 :             : writepos_m(0)
      15           26 :             , readpos_m(0)
      16           26 :             , buffer_m("buffer", size) {}
      17              : 
      18              :         template <class... Properties>
      19              :         template <typename T, class... ViewArgs>
      20            0 :         void Archive<Properties...>::serialize(const Kokkos::View<T*, ViewArgs...>& view,
      21              :                                                size_type nsends) {
      22              :             using exec_space  = typename Kokkos::View<T*, ViewArgs...>::execution_space;
      23              :             using policy_type = Kokkos::RangePolicy<exec_space>;
      24              : 
      25            0 :             size_t size = sizeof(T);
      26            0 :             Kokkos::parallel_for(
      27            0 :                 "Archive::serialize()", policy_type(0, nsends),
      28            0 :                 KOKKOS_CLASS_LAMBDA(const size_type i) {
      29            0 :                     std::memcpy(buffer_m.data() + i * size + writepos_m, view.data() + i, size);
      30              :                 });
      31            0 :             Kokkos::fence();
      32            0 :             writepos_m += size * nsends;
      33            0 :         }
      34              : 
      35              :         template <class... Properties>
      36              :         template <typename T, unsigned Dim, class... ViewArgs>
      37            0 :         void Archive<Properties...>::serialize(
      38              :             const Kokkos::View<Vector<T, Dim>*, ViewArgs...>& view, size_type nsends) {
      39              :             using exec_space = typename Kokkos::View<T*, ViewArgs...>::execution_space;
      40              : 
      41            0 :             size_t size = sizeof(T);
      42              :             // Default index type for range policies is int64,
      43              :             // so we have to explicitly specify size_type (uint64)
      44              :             using mdrange_t =
      45              :                 Kokkos::MDRangePolicy<Kokkos::Rank<2>, Kokkos::IndexType<size_type>, exec_space>;
      46            0 :             Kokkos::parallel_for(
      47              :                 "Archive::serialize()",
      48              :                 // The constructor for Kokkos range policies always
      49              :                 // expects int64 regardless of index type provided
      50              :                 // by template parameters, so the typecast is necessary
      51              :                 // to avoid compiler warnings
      52            0 :                 mdrange_t({0, 0}, {(long int)nsends, Dim}),
      53            0 :                 KOKKOS_CLASS_LAMBDA(const size_type i, const size_t d) {
      54            0 :                     std::memcpy(buffer_m.data() + (Dim * i + d) * size + writepos_m,
      55            0 :                                 &(*(view.data() + i))[d], size);
      56              :                 });
      57            0 :             Kokkos::fence();
      58            0 :             writepos_m += Dim * size * nsends;
      59            0 :         }
      60              : 
      61              :         template <class... Properties>
      62              :         template <typename T, class... ViewArgs>
      63            0 :         void Archive<Properties...>::deserialize(Kokkos::View<T*, ViewArgs...>& view,
      64              :                                                  size_type nrecvs) {
      65              :             using exec_space  = typename Kokkos::View<T*, ViewArgs...>::execution_space;
      66              :             using policy_type = Kokkos::RangePolicy<exec_space>;
      67              : 
      68            0 :             size_t size = sizeof(T);
      69            0 :             if (nrecvs > view.extent(0)) {
      70            0 :                 Kokkos::realloc(view, nrecvs);
      71              :             }
      72            0 :             Kokkos::parallel_for(
      73            0 :                 "Archive::deserialize()", policy_type(0, nrecvs),
      74            0 :                 KOKKOS_CLASS_LAMBDA(const size_type i) {
      75            0 :                     std::memcpy(view.data() + i, buffer_m.data() + i * size + readpos_m, size);
      76              :                 });
      77              :             // Wait for deserialization kernel to complete
      78              :             // (as with serialization kernels)
      79            0 :             Kokkos::fence();
      80            0 :             readpos_m += size * nrecvs;
      81            0 :         }
      82              : 
      83              :         template <class... Properties>
      84              :         template <typename T, unsigned Dim, class... ViewArgs>
      85            0 :         void Archive<Properties...>::deserialize(Kokkos::View<Vector<T, Dim>*, ViewArgs...>& view,
      86              :                                                  size_type nrecvs) {
      87              :             using exec_space = typename Kokkos::View<T*, ViewArgs...>::execution_space;
      88              : 
      89            0 :             size_t size = sizeof(T);
      90            0 :             if (nrecvs > view.extent(0)) {
      91            0 :                 Kokkos::realloc(view, nrecvs);
      92              :             }
      93              :             using mdrange_t =
      94              :                 Kokkos::MDRangePolicy<Kokkos::Rank<2>, Kokkos::IndexType<size_type>, exec_space>;
      95            0 :             Kokkos::parallel_for(
      96            0 :                 "Archive::deserialize()", mdrange_t({0, 0}, {(long int)nrecvs, Dim}),
      97            0 :                 KOKKOS_CLASS_LAMBDA(const size_type i, const size_t d) {
      98            0 :                     std::memcpy(&(*(view.data() + i))[d],
      99            0 :                                 buffer_m.data() + (Dim * i + d) * size + readpos_m, size);
     100              :                 });
     101            0 :             Kokkos::fence();
     102            0 :             readpos_m += Dim * size * nrecvs;
     103            0 :         }
     104              :     }  // namespace detail
     105              : }  // namespace ippl