Non-contiguous data handling, restructure for multiple transports

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));
+  std::vector<KokkosComm::Req<>> reqs = 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();
-  }
+  KokkosComm::wait_all(reqs);
 }
 
 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,21 @@
 
 #pragma once
 
+#include "KokkosComm_fwd.hpp"
+
+// transport declarations
+// TODO: could probably be moved to a per-transport file to be included
+#if defined(KOKKOSCOMM_TRANSPORT_MPI)
+#include "mpi/KokkosComm_mpi.hpp"
+#include "mpi/KokkosComm_mpi_handle.hpp"
+#include "mpi/KokkosComm_mpi_plan.hpp"
+#include "mpi/KokkosComm_mpi_isend.hpp"
+#include "mpi/KokkosComm_mpi_irecv.hpp"
+#else
+#error at least one transport must be defined
+#endif
+
+#include "KokkosComm_point_to_point.hpp"
 #include "KokkosComm_collective.hpp"
 #include "KokkosComm_version.hpp"
 #include "KokkosComm_isend.hpp"
@@ -27,14 +42,10 @@
 #include "KokkosComm_concepts.hpp"
 #include "KokkosComm_comm_mode.hpp"
 
-#include <Kokkos_Core.hpp>
-
 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,19 @@
+//@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
+
+#define KOKKOSCOMM_TRANSPORT_MPI  // TODO: set through CMake

src/KokkosComm_fwd.hpp

@@ -0,0 +1,54 @@
+//@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 <vector>
+
+#include "KokkosComm_concepts.hpp"
+#include "KokkosComm_config.hpp"
+
+namespace KokkosComm {
+#if defined(KOKKOSCOMM_TRANSPORT_MPI)
+class Mpi;
+using DefaultTransport  = Mpi;
+using FallbackTransport = Mpi;
+#else
+#error at least one transport must be defined
+#endif
+
+template <Transport TRANSPORT = DefaultTransport>
+class Req;
+
+template <KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace, Transport TRANSPORT = DefaultTransport>
+class Handle;
+
+namespace Impl {
+
+template <Dispatch TDispatch, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+          Transport TRANSPORT = DefaultTransport>
+struct Plan;
+
+template <KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+          Transport TRANSPORT = DefaultTransport>
+struct Irecv;
+template <KokkosView SendView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+          Transport TRANSPORT = DefaultTransport>
+struct Isend;
+
+}  // namespace Impl
+
+}  // namespace KokkosComm

src/KokkosComm_plan.hpp

@@ -0,0 +1,39 @@
+//@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 <Kokkos_Core.hpp>
+
+#include "KokkosComm_fwd.hpp"
+
+namespace KokkosComm {
+
+template <Dispatch TDispatch, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+          Transport TRANSPORT = DefaultTransport>
+std::vector<Req<TRANSPORT>> plan(Handle<ExecSpace, TRANSPORT> &handle, TDispatch d) {
+  return Impl::Plan<TDispatch, ExecSpace, TRANSPORT>(handle, d).reqs;
+}
+
+template <Dispatch TDispatch, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+          Transport TRANSPORT = DefaultTransport>
+std::vector<Req<TRANSPORT>> plan(const ExecSpace &space, MPI_Comm comm, TDispatch d) {
+  Handle<ExecSpace, TRANSPORT> handle(space, comm);
+  auto ret = plan<TDispatch, ExecSpace, TRANSPORT>(handle, d);
+  return ret;
+}
+
+}  // namespace KokkosComm

src/KokkosComm_point_to_point.hpp

@@ -0,0 +1,59 @@
+//@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 <Kokkos_Core.hpp>
+
+#include "KokkosComm_fwd.hpp"
+#include "KokkosComm_concepts.hpp"
+#include "KokkosComm_plan.hpp"
+
+namespace KokkosComm {
+
+template <KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+          Transport TRANSPORT = DefaultTransport>
+void irecv(Handle<ExecSpace, TRANSPORT> &h, RecvView &rv, int src, int tag) {
+  Impl::Irecv<RecvView, ExecSpace, TRANSPORT>(h, rv, src, tag);
+}
+
+template <KokkosView SendView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+          Transport TRANSPORT = DefaultTransport>
+void isend(Handle<ExecSpace, TRANSPORT> &h, SendView &sv, int dest, int tag) {
+  Impl::Isend<SendView, ExecSpace, TRANSPORT>(h, sv, dest, tag);
+}
+
+// TODO: can these go in MPI somewhere?
+#if defined(KOKKOSCOMM_TRANSPORT_MPI)
+template <KokkosView SendView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace>
+Req<Mpi> isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm) {
+  auto reqs =
+      KokkosComm::plan(space, comm, [=](Handle<ExecSpace, Mpi> &handle) { KokkosComm::isend(handle, sv, dest, tag); });
+  assert(reqs.size() == 1 && "Internal KokkosComm developer error");
+  return reqs[0];
+}
+
+template <KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace>
+Req<Mpi> irecv(const ExecSpace &space, const RecvView &rv, int dest, int tag, MPI_Comm comm) {
+  auto reqs =
+      KokkosComm::plan(space, comm, [=](Handle<ExecSpace, Mpi> &handle) { KokkosComm::irecv(handle, rv, dest, tag); });
+  assert(reqs.size() == 1 && "Internal KokkosComm developer error");
+  return reqs[0];
+}
+
+#endif
+
+}  // 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_concepts.hpp

@@ -16,14 +16,30 @@
 
 #pragma once
 
+#include <type_traits>
+
 #include <Kokkos_Core.hpp>
 
 namespace KokkosComm {
 
+namespace Impl {
+
+// fallback - most types are not a KokkosComm transport
+template <typename T>
+struct is_transport : public std::false_type {};
+}  // namespace Impl
+
 template <typename T>
 concept KokkosView = Kokkos::is_view_v<T>;
 
 template <typename T>
 concept KokkosExecutionSpace = Kokkos::is_execution_space_v<T>;
 
+// TODO: a function that takes something convertible to a handle and returns void
+template <typename T>
+concept Dispatch = true;
+
+template <typename T>
+concept Transport = KokkosComm::Impl::is_transport<T>::value;
+
 }  // namespace KokkosComm

src/impl/KokkosComm_contiguous.hpp

@@ -0,0 +1,63 @@
+//@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"
+#include "KokkosComm_traits.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_irecv.hpp

@@ -21,7 +21,6 @@
 #include <Kokkos_Core.hpp>
 
 #include "KokkosComm_pack_traits.hpp"
-#include "KokkosComm_request.hpp"
 #include "KokkosComm_traits.hpp"
 
 // impl
@@ -45,9 +44,9 @@ void irecv(RecvView &rv, int src, int tag, MPI_Comm comm, MPI_Request &req) {
 }
 
 template <KokkosView RecvView>
-KokkosComm::Req irecv(RecvView &rv, int src, int tag, MPI_Comm comm) {
+KokkosComm::Req<KokkosComm::Mpi> irecv(RecvView &rv, int src, int tag, MPI_Comm comm) {
   Kokkos::Tools::pushRegion("KokkosComm::Impl::irecv");
-  KokkosComm::Req req;
+  KokkosComm::Req<KokkosComm::Mpi> req;
   irecv(rv, src, tag, comm, req.mpi_req());
   return req;
 }