MpiDataype packer

src/KokkosComm_pack_traits.hpp

@@ -37,7 +37,8 @@ struct PackTraits {
 /*! \brief This can be specialized to do custom behavior for a particular view*/
 template <KokkosView View>
 struct PackTraits<View> {
-  using packer_type = Impl::Packer::DeepCopy<View>;
+  // using packer_type = Impl::Packer::DeepCopy<View>;
+  using packer_type = Impl::Packer::MpiDatatype<View>;
 
   static bool needs_unpack(const View &v) {
     return !Traits<View>::is_contiguous(v);

src/impl/KokkosComm_packer.hpp

@@ -17,6 +17,7 @@
 #pragma once
 
 #include "KokkosComm_concepts.hpp"
+#include "KokkosComm_types.hpp"
 #include "KokkosComm_include_mpi.hpp"
 
 namespace KokkosComm::Impl {
@@ -82,5 +83,46 @@ struct DeepCopy {
   }
 };
 
+template <KokkosView View>
+struct MpiDatatype {
+  using non_const_packed_view_type = View;
+  using args_type                  = MpiArgs<non_const_packed_view_type>;
+
+  // don't actually allocate - return the provided view, but with
+  // a datatype that describes the data in the view
+  template <typename ExecSpace>
+  static args_type allocate_packed_for(const ExecSpace &space,
+                                       const std::string & /*label*/,
+                                       const View &src) {
+    using ValueType = typename View::value_type;
+
+    MPI_Datatype type = mpi_type<ValueType>();
+    for (size_t d = 0; d < KokkosComm::Traits<View>::rank(); ++d) {
+      MPI_Datatype newtype;
+      MPI_Type_create_hvector(src.extent(d) /*count*/, 1 /*block length*/,
+                              src.stride(d) * sizeof(ValueType), type,
+                              &newtype);
+      type = newtype;
+    }
+    MPI_Type_commit(&type);
+    MPI_Aint lb, extent;
+    MPI_Type_get_extent(type, &lb, &extent);
+    return args_type(src, type, 1);
+  }
+
+  // pack is a no-op: rely on MPI's datatype engine
+  template <typename ExecSpace>
+  static args_type pack(const ExecSpace &space, const View &src) {
+    return allocate_packed_for(space, "", src);
+  }
+
+  // unpack is a no-op: rely on MPI's datatype engine
+  template <typename ExecSpace>
+  static void unpack_into(const ExecSpace &space, const View &dst,
+                          const non_const_packed_view_type &src) {
+    return;
+  }
+};
+
 }  // namespace Packer
 }  // namespace KokkosComm::Impl