Initial attempt at comm/space wrapping and noncontig views

Author: cwpearson

perf_tests/test_2dhalo.cpp

@@ -46,22 +46,20 @@ void send_recv(benchmark::State &, MPI_Comm comm, const Space &space, int nx, in
   auto ym1_s = Kokkos::subview(v, make_pair(1, nx + 1), 1, Kokkos::ALL);
   auto ym1_r = Kokkos::subview(v, make_pair(1, nx + 1), 0, Kokkos::ALL);
 
-  std::vector<KokkosComm::Req> reqs;
-  // std::cerr << get_rank(rx, ry) << " -> " << get_rank(xp1, ry) << "\n";
-  reqs.push_back(KokkosComm::isend(space, xp1_s, get_rank(xp1, ry), 0, comm));
-  reqs.push_back(KokkosComm::isend(space, xm1_s, get_rank(xm1, ry), 1, comm));
-  reqs.push_back(KokkosComm::isend(space, yp1_s, get_rank(rx, yp1), 2, comm));
-  reqs.push_back(KokkosComm::isend(space, ym1_s, get_rank(rx, ym1), 3, comm));
+  KokkosComm::Handle<Space> h = KokkosComm::plan(space, comm, [=](KokkosComm::Handle<Space> &handle) {
+    KokkosComm::isend(handle, xp1_s, get_rank(xp1, ry), 0);
+    KokkosComm::isend(handle, xm1_s, get_rank(xm1, ry), 1);
+    KokkosComm::isend(handle, yp1_s, get_rank(rx, yp1), 2);
+    KokkosComm::isend(handle, ym1_s, get_rank(rx, ym1), 3);
+  });
 
   KokkosComm::recv(space, xm1_r, get_rank(xm1, ry), 0, comm);
   KokkosComm::recv(space, xp1_r, get_rank(xp1, ry), 1, comm);
   KokkosComm::recv(space, ym1_r, get_rank(rx, ym1), 2, comm);
   KokkosComm::recv(space, yp1_r, get_rank(rx, yp1), 3, comm);
 
   // wait for comm
-  for (KokkosComm::Req &req : reqs) {
-    req.wait();
-  }
+  h.wait();
 }
 
 void benchmark_2dhalo(benchmark::State &state) {

src/CMakeLists.txt

@@ -6,6 +6,7 @@ target_include_directories(KokkosComm INTERFACE
 )
 target_include_directories(KokkosComm INTERFACE
     $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}/impl>
+    $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}/mpi>
 )
 target_include_directories(KokkosComm INTERFACE
     $<INSTALL_INTERFACE:include>

src/KokkosComm.hpp

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include "KokkosComm_point_to_point.hpp"
 #include "KokkosComm_collective.hpp"
 #include "KokkosComm_version.hpp"
 #include "KokkosComm_isend.hpp"
@@ -33,8 +34,6 @@ namespace KokkosComm {
 
 using Impl::alltoall;
 using Impl::barrier;
-using Impl::irecv;
-using Impl::isend;
 using Impl::recv;
 using Impl::send;
 

src/KokkosComm_config.hpp

@@ -0,0 +1,29 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#pragma once
+
+#include "mpi/KokkosComm_mpi.hpp"
+#include "KokkosComm_concepts.hpp"
+
+namespace KokkosComm {
+
+using GenericTransport = ::KokkosComm::Mpi;
+using SpecialTransport = ::KokkosComm::Mpi;
+template <KokkosExecutionSpace ExecSpace>
+using Handle = Mpi::Handle<ExecSpace>;
+
+}  // namespace KokkosComm

src/KokkosComm_plan.hpp

@@ -0,0 +1,26 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#pragma once
+
+namespace KokkosComm {
+
+template <KokkosExecutionSpace ExecSpace, typename CommFunc>
+Mpi::Handle<ExecSpace> plan(const ExecSpace &space, MPI_Comm comm, CommFunc f) {
+  return Mpi::Plan<ExecSpace, CommFunc>(space, comm, f).handle();
+}
+
+}  // namespace KokkosComm

src/KokkosComm_point_to_point.hpp

@@ -0,0 +1,56 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#pragma once
+
+#include "KokkosComm_api.hpp"
+#include "KokkosComm_concepts.hpp"
+#include "KokkosComm_config.hpp"
+#include "KokkosComm_plan.hpp"
+
+namespace KokkosComm {
+
+template <typename Handle, KokkosView RecvView>
+void irecv(Handle &h, RecvView &rv, int src, int tag) {
+  if constexpr (Impl::api_avail_v<SpecialTransport, Impl::Api::Irecv>) {
+    SpecialTransport::irecv(h, rv, src, tag);
+  } else {
+    GenericTransport::irecv(h, rv, src, tag);
+  }
+}
+
+template <KokkosExecutionSpace ExecSpace, KokkosView RecvView>
+KokkosComm::Handle<ExecSpace> irecv(const ExecSpace &space, const RecvView &rv, int dest, int tag, MPI_Comm comm) {
+  using MyHandle = KokkosComm::Handle<ExecSpace>;
+  return KokkosComm::plan(space, comm, [=](MyHandle &handle) { KokkosComm::irecv(handle, rv, dest, tag); });
+}
+
+template <typename Handle, KokkosView SendView>
+void isend(Handle &h, SendView &sv, int src, int tag) {
+  if constexpr (Impl::api_avail_v<SpecialTransport, Impl::Api::Isend>) {
+    SpecialTransport::isend(h, sv, src, tag);
+  } else {
+    GenericTransport::isend(h, sv, src, tag);
+  }
+}
+
+template <KokkosExecutionSpace ExecSpace, KokkosView SendView>
+KokkosComm::Handle<ExecSpace> isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm) {
+  using MyHandle = KokkosComm::Handle<ExecSpace>;
+  return KokkosComm::plan(space, comm, [=](MyHandle &handle) { KokkosComm::isend(handle, sv, dest, tag); });
+}
+
+}  // namespace KokkosComm

src/impl/KokkosComm_api.hpp

@@ -0,0 +1,30 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#pragma once
+
+namespace KokkosComm::Impl {
+
+enum class Api { Irecv, Isend };
+
+// catch-all: no transports implement any APIs
+template <typename Transport, Impl::Api API>
+struct api_avail : public std::false_type {};
+
+template <typename Transport, Impl::Api API>
+constexpr bool api_avail_v = api_avail<Transport, API>::value;
+
+}  // namespace KokkosComm::Impl

src/impl/KokkosComm_contiguous.hpp

@@ -0,0 +1,62 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#pragma once
+
+#include <string>
+
+#include <Kokkos_Core.hpp>
+
+#include "KokkosComm_concepts.hpp"
+
+namespace KokkosComm::Impl {
+
+template <KokkosView View>
+struct contiguous_view {
+  using type = Kokkos::View<typename View::non_const_data_type, Kokkos::LayoutRight, typename View::memory_space>;
+};
+
+template <KokkosView View>
+using contiguous_view_t = contiguous_view<View>::type;
+
+template <KokkosView View, KokkosExecutionSpace Space>
+auto allocate_contiguous_for(const Space &space, const std::string &label, View &v) {
+  using non_const_packed_view_type = contiguous_view_t<View>;
+
+  if constexpr (KokkosComm::rank<View>() == 1) {
+    return non_const_packed_view_type(Kokkos::view_alloc(space, Kokkos::WithoutInitializing, label), v.extent(0));
+  } else if constexpr (KokkosComm::rank<View>() == 2) {
+    return non_const_packed_view_type(Kokkos::view_alloc(space, Kokkos::WithoutInitializing, label), v.extent(0),
+                                      v.extent(1));
+  } else {
+    static_assert(std::is_void_v<View>, "allocate_contiguous_for for views > rank 2 not implemented");
+  }
+}
+
+template <KokkosExecutionSpace Space, KokkosView DstView, KokkosView SrcView>
+auto resize_contiguous_for(const Space &space, DstView &out, const SrcView &in) {
+  static_assert(DstView::rank == SrcView::rank, "");
+
+  if constexpr (KokkosComm::rank<DstView>() == 1) {
+    Kokkos::realloc(Kokkos::view_alloc(space, Kokkos::WithoutInitializing), out, in.extent(0));
+  } else if constexpr (KokkosComm::rank<DstView>() == 2) {
+    Kokkos::realloc(Kokkos::view_alloc(space, Kokkos::WithoutInitializing), out, in.extent(0), in.extent(1));
+  } else {
+    static_assert(std::is_void_v<DstView>, "realloc_contiguous_for for views > rank 2 not implemented");
+  }
+}
+
+}  // namespace KokkosComm::Impl

src/impl/KokkosComm_types.hpp

@@ -107,4 +107,50 @@ MPI_Datatype mpi_type() {
 
 template <typename Scalar>
 inline MPI_Datatype mpi_type_v = mpi_type<Scalar>();
+
+template <KokkosView View>
+MPI_Datatype view_mpi_type(const View &view) {
+#define USE_CACHE
+
+#if defined(USE_CACHE)
+  constexpr int P = 2 * View::rank;
+  using Key       = std::array<int, P>;
+
+  auto key_from = [](const View &v) -> Key {
+    Key key;
+    for (size_t d = 0; d < View::rank; d++) {
+      key[d]     = v.extent(d);
+      key[d + 1] = v.stride(d);
+    }
+    return key;
+  };
+
+  static std::map<Key, MPI_Datatype> cache;
+
+  Key key = key_from(view);
+  if (cache.count(key) > 0) {
+    return cache[key];
+  }
+#endif
+
+  using value_type  = typename View::non_const_value_type;
+  MPI_Datatype type = mpi_type_v<value_type>;
+
+  // This doesn't work for 1D contiguous views into reduce because it
+  // represents the whole 1D view as 1 Hvector, rather than N elements.
+  // FIXME: is there a more generic way to handle this, maybe by treating
+  // the last dimension specially under certain circumstances?
+  for (size_t d = 0; d < KokkosComm::rank<View>(); ++d) {
+    MPI_Datatype newtype;
+    MPI_Type_create_hvector(KokkosComm::extent(view, d) /*count*/, 1 /*block length*/,
+                            KokkosComm::stride(view, d) * sizeof(value_type), type, &newtype);
+    type = newtype;
+  }
+  MPI_Type_commit(&type);
+#if defined(USE_CACHE)
+  cache[key] = type;
+#endif
+  return type;
+}
+
 };  // namespace KokkosComm::Impl