has_serialize.hpp

#ifndef __COMMON_HAS_SERIALIZE_HPP__
#define __COMMON_HAS_SERIALIZE_HPP__
 
#include <Kokkos_Core.hpp>
#include <biocma_cst_config.hpp>
#include <type_traits>
#include <vector>
 
#define IF_ENABLE_CEREAL_SERDE                                                 \
  typename std::enable_if<                                                     \
      AutoGenerated::FlagCompileTime::use_cereal_serde>::type
 
namespace
{
  template <typename, typename T> struct __has_serialize
  {
    static_assert(std::integral_constant<T, false>::value,
                  "Second template parameter needs to be of function type.");
  };
 
  // specialization that does the checking
 
  template <typename C, typename Ret, typename... Args>
  struct __has_serialize<C, Ret(Args...)>
  {
  private:
    template <typename T>
    static constexpr auto check(T*) -> typename std::is_same<
        decltype(std::declval<T>().serialize(std::declval<Args>()...)),
        Ret>::type;
 
    template <typename> static constexpr std::false_type check(...);
 
    using type = decltype(check<C>(0));
 
  public:
    static constexpr bool value = type::value;
  };
 
  template <class ViewType, typename Enable = void> struct ViewSerialization;
 
  // Specialization for 1D Kokkos::View
  template <class ViewType>
  struct ViewSerialization<
      ViewType,
      typename std::enable_if<
          ViewType::rank == 1
          && AutoGenerated::FlagCompileTime::use_cereal_serde>::type>
  {
    template <class Archive>
    static void
    serialize(Archive& ar, const ViewType& view, std::size_t n_used_elements)
    {
      using T = typename ViewType::value_type;
      auto temp_host = Kokkos::create_mirror_view_and_copy(
          Kokkos::DefaultHostExecutionSpace(), view);
      std::vector<T> data_vector(temp_host.data(),
                                 temp_host.data() + n_used_elements);
      ar(data_vector);
    }
 
    template <class Archive>
    static void
    deserialize(Archive& ar, ViewType& view)
    {
      using T = typename ViewType::value_type;
 
      std::vector<T> data_vector;
      ar(data_vector);
 
      view = ViewType("deser", data_vector.size());
 
      auto tmpdata = ViewType(data_vector.data(), data_vector.size());
      Kokkos::deep_copy(view, tmpdata);
    }
  };
 
  // Specialization for 2D Kokkos::View
  template <class ViewType>
  struct ViewSerialization<
      ViewType,
      typename std::enable_if<
          ViewType::rank == 2
          && AutoGenerated::FlagCompileTime::use_cereal_serde>::type>
  {
    template <class Archive>
    static void
    serialize(Archive& ar, const ViewType& view, std::size_t n_used_elements)
    {
      (void)n_used_elements;
      using T = typename ViewType::value_type;
      auto temp_host = Kokkos::create_mirror_view_and_copy(
          Kokkos::DefaultHostExecutionSpace(), view);
      const auto rows = view.extent(0);
      const auto cols = view.extent(1);
      ar(rows);
      ar(cols);
      std::vector<T> data_vector(temp_host.data(),
                                 temp_host.data() + rows * cols);
      ar(data_vector);
    }
 
    template <class Archive>
    static void
    deserialize(Archive& ar, ViewType& view)
    {
      using T = typename ViewType::value_type;
      using L = typename ViewType::array_layout;
      std::size_t rows{};
      std::size_t cols{};
      std::vector<T> data_vector;
      ar(rows, cols);
      ar(data_vector);
      view = ViewType("deser", rows, cols);
      auto tmpdata = ViewType(data_vector.data(), L(rows, cols));
 
      Kokkos::deep_copy(view, tmpdata);
    }
  };
} // namespace
 
// Helper to check if type can be serialized
template <class C, typename T>
consteval bool
has_serialize()
{
  return __has_serialize<C, void(T&)>::value;
}
 
// Helper to serialize Kokkos View in any memory space
template <class Archive, class ViewType>
void
serialize_view(Archive& archive, const ViewType& view)
{
  // using T = typename ViewType::value_type;
  ViewSerialization<ViewType>::serialize(archive, view, view.extent(0));
}
 
// Helper to deserialize Kokkos View in any memory space
template <class Archive, class ViewType>
void
deserialize_view(Archive& archive, ViewType& view)
{
  ViewSerialization<ViewType>::deserialize(archive, view);
}
 
#endif