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        13052 :         Archive<Properties...>::Archive(size_type size)
      14        13052 :             : writepos_m(0)
      15        13052 :             , readpos_m(0)
      16        13052 :             , buffer_m("buffer", size) {}
      17              : 
      18              :         template <class... Properties>
      19              :         template <typename T, class... ViewArgs>
      20         9676 :         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         9676 :             size_t size = sizeof(T);
      26         9676 :             Kokkos::parallel_for(
      27         9676 :                 "Archive::serialize()", policy_type(0, nsends),
      28      1994404 :                 KOKKOS_CLASS_LAMBDA(const size_type i) {
      29       388644 :                     std::memcpy(buffer_m.data() + i * size + writepos_m, view.data() + i, size);
      30              :                 });
      31         9676 :             Kokkos::fence();
      32         9676 :             writepos_m += size * nsends;
      33         9676 :         }
      34              : 
      35              :         template <class... Properties>
      36              :         template <typename T, unsigned Dim, class... ViewArgs>
      37          212 :         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          212 :             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          212 :             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          212 :                 mdrange_t({0, 0}, {(long int)nsends, Dim}),
      53       393092 :                 KOKKOS_CLASS_LAMBDA(const size_type i, const size_t d) {
      54       196334 :                     std::memcpy(buffer_m.data() + (Dim * i + d) * size + writepos_m,
      55       196334 :                                 &(*(view.data() + i))[d], size);
      56              :                 });
      57          212 :             Kokkos::fence();
      58          212 :             writepos_m += Dim * size * nsends;
      59          212 :         }
      60              : 
      61              :         template <class... Properties>
      62              :         template <typename T, class... ViewArgs>
      63         9676 :         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         9676 :             size_t size = sizeof(T);
      69         9676 :             if (nrecvs > view.extent(0)) {
      70           72 :                 Kokkos::realloc(view, nrecvs);
      71              :             }
      72         9676 :             Kokkos::parallel_for(
      73         9676 :                 "Archive::deserialize()", policy_type(0, nrecvs),
      74      1994404 :                 KOKKOS_CLASS_LAMBDA(const size_type i) {
      75       388644 :                     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         9676 :             Kokkos::fence();
      80         9676 :             readpos_m += size * nrecvs;
      81         9676 :         }
      82              : 
      83              :         template <class... Properties>
      84              :         template <typename T, unsigned Dim, class... ViewArgs>
      85          212 :         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          212 :             size_t size = sizeof(T);
      90          212 :             if (nrecvs > view.extent(0)) {
      91           60 :                 Kokkos::realloc(view, nrecvs);
      92              :             }
      93              :             using mdrange_t =
      94              :                 Kokkos::MDRangePolicy<Kokkos::Rank<2>, Kokkos::IndexType<size_type>, exec_space>;
      95          212 :             Kokkos::parallel_for(
      96          212 :                 "Archive::deserialize()", mdrange_t({0, 0}, {(long int)nrecvs, Dim}),
      97       393092 :                 KOKKOS_CLASS_LAMBDA(const size_type i, const size_t d) {
      98       196334 :                     std::memcpy(&(*(view.data() + i))[d],
      99       196334 :                                 buffer_m.data() + (Dim * i + d) * size + readpos_m, size);
     100              :                 });
     101          212 :             Kokkos::fence();
     102          212 :             readpos_m += Dim * size * nrecvs;
     103          212 :         }
     104              :     }  // namespace detail
     105              : }  // namespace ippl