rework to support waitall

cwpearson · cwpearson · commit c1f1b67ab5e7 · 2024-06-21T15:32:05.000-06:00
diff --git a/perf_tests/test_2dhalo.cpp b/perf_tests/test_2dhalo.cpp
@@ -46,7 +46,7 @@ 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);
 
-  KokkosComm::Handle<Space> h = KokkosComm::plan(space, comm, [=](KokkosComm::Handle<Space> &handle) {
+  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);
@@ -59,7 +59,7 @@ void send_recv(benchmark::State &, MPI_Comm comm, const Space &space, int nx, in
   KokkosComm::recv(space, yp1_r, get_rank(rx, yp1), 3, comm);
 
   // wait for comm
-  KokkosComm::wait(h);
+  KokkosComm::wait_all(reqs);
 }
 
 void benchmark_2dhalo(benchmark::State &state) {
diff --git a/src/KokkosComm_fwd.hpp b/src/KokkosComm_fwd.hpp
@@ -16,6 +16,8 @@
 
 #pragma once
 
+#include <vector>
+
 #include "KokkosComm_concepts.hpp"
 #include "KokkosComm_config.hpp"
 
@@ -28,15 +30,18 @@ using FallbackTransport = Mpi;
 #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;
 
-template <Dispatch DISPATCH, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
-          Transport TRANSPORT = DefaultTransport>
-class Plan;
-
 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;
diff --git a/src/KokkosComm_plan.hpp b/src/KokkosComm_plan.hpp
@@ -22,15 +22,18 @@
 
 namespace KokkosComm {
 
-template <Dispatch DISPATCH, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+template <Dispatch TDispatch, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
           Transport TRANSPORT = DefaultTransport>
-Handle<ExecSpace, TRANSPORT> plan(const ExecSpace &space, MPI_Comm comm, DISPATCH d) {
-  return Plan<DISPATCH, ExecSpace, TRANSPORT>(space, comm, d).handle();
+std::vector<Req<TRANSPORT>> plan(Handle<ExecSpace, TRANSPORT> &handle, TDispatch d) {
+  return Impl::Plan<TDispatch, ExecSpace, TRANSPORT>(handle, d).reqs;
 }
 
-template <Dispatch DISPATCH, KokkosExecutionSpace ExecSpace, Transport TRANSPORT>
-void plan(Handle<ExecSpace, TRANSPORT> &handle, DISPATCH d) {
-  Plan<DISPATCH, ExecSpace, TRANSPORT>(handle, d);
+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
diff --git a/src/KokkosComm_point_to_point.hpp b/src/KokkosComm_point_to_point.hpp
@@ -39,15 +39,19 @@ void isend(Handle<ExecSpace, TRANSPORT> &h, SendView &sv, int dest, int tag) {
 // TODO: can these go in MPI somewhere?
 #if defined(KOKKOSCOMM_TRANSPORT_MPI)
 template <KokkosView SendView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace>
-KokkosComm::Handle<ExecSpace, Mpi> isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm) {
-  return KokkosComm::plan(space, comm,
-                          [=](Handle<ExecSpace, Mpi> &handle) { KokkosComm::isend(handle, sv, dest, tag); });
+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>
-KokkosComm::Handle<ExecSpace, Mpi> irecv(const ExecSpace &space, const RecvView &rv, int dest, int tag, MPI_Comm comm) {
-  return KokkosComm::plan(space, comm,
-                          [=](Handle<ExecSpace, Mpi> &handle) { KokkosComm::irecv(handle, rv, dest, tag); });
+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
diff --git a/src/impl/KokkosComm_irecv.hpp b/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;
 }
diff --git a/src/impl/KokkosComm_isend.hpp b/src/impl/KokkosComm_isend.hpp
@@ -22,7 +22,6 @@
 
 #include "KokkosComm_concepts.hpp"
 #include "KokkosComm_pack_traits.hpp"
-#include "KokkosComm_request.hpp"
 #include "KokkosComm_traits.hpp"
 #include "KokkosComm_comm_mode.hpp"
 
@@ -46,10 +45,10 @@ void isend(const SendView &sv, int dest, int tag, MPI_Comm comm, MPI_Request &re
 }
 
 template <CommMode SendMode = CommMode::Default, KokkosExecutionSpace ExecSpace, KokkosView SendView>
-KokkosComm::Req isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm) {
+KokkosComm::Req<KokkosComm::Mpi> isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm) {
   Kokkos::Tools::pushRegion("KokkosComm::Impl::isend");
 
-  KokkosComm::Req req;
+  KokkosComm::Req<KokkosComm::Mpi> req;
 
   using KCT  = KokkosComm::Traits<SendView>;
   using KCPT = KokkosComm::PackTraits<SendView>;
diff --git a/src/mpi/KokkosComm_mpi_handle.hpp b/src/mpi/KokkosComm_mpi_handle.hpp
@@ -18,6 +18,8 @@
 
 #include "KokkosComm_fwd.hpp"
 
+#include "KokkosComm_mpi_req.hpp"
+
 namespace KokkosComm {
 
 /*
@@ -35,41 +37,23 @@ class Handle<ExecSpace, Mpi> {
  public:
   using execution_space = ExecSpace;
 
-  // a type-erased view. Request uses these to keep temporary views alive for
-  // the lifetime of "Immediate" MPI operations
-  struct ViewHolderBase {
-    virtual ~ViewHolderBase() {}
-  };
-  template <typename V>
-  struct ViewHolder : ViewHolderBase {
-    ViewHolder(const V &v) : v_(v) {}
-    V v_;
-  };
-
   Handle(const execution_space &space, MPI_Comm comm) : space_(space), comm_(comm), preCommFence_(false) {}
 
-  // template <KokkosView View>
-  // void impl_track_view(const View &v) {
-  //     views_.push_back(std::make_shared<ViewHolder<View>>(v));
-  // }
+  MPI_Comm &mpi_comm() { return comm_; }
+  const execution_space &space() const { return space_; }
 
   void impl_add_pre_comm_fence() { preCommFence_ = true; }
 
-  void impl_track_mpi_request(MPI_Request req) { reqs_.push_back(req); }
-
   void impl_add_alloc(std::function<void()> f) { allocs_.push_back(f); }
 
   void impl_add_pre_copy(std::function<void()> f) { preCopies_.push_back(f); }
 
   void impl_add_comm(std::function<void()> f) { comms_.push_back(f); }
 
-  void impl_add_post_wait(std::function<void()> f) { postWaits_.push_back(f); }
-
-  MPI_Comm &mpi_comm() { return comm_; }
-
-  const execution_space &space() const { return space_; }
+  void impl_track_req(const Req<Mpi> &req) { reqs_.push_back(req); }
 
   void impl_run() {
+    std::cerr << __FILE__ << ":" << __LINE__ << " impl_run\n";
     for (const auto &f : allocs_) f();
     for (const auto &f : preCopies_) f();
     if (preCommFence_) {
@@ -82,10 +66,9 @@ class Handle<ExecSpace, Mpi> {
     comms_.clear();
   }
 
- private:
-  template <KokkosExecutionSpace ES>
-  friend void wait(Handle<ES, Mpi> &handle);
+  std::vector<Req<Mpi>> &impl_reqs() { return reqs_; }
 
+ private:
   execution_space space_;
   MPI_Comm comm_;
 
@@ -96,23 +79,8 @@ class Handle<ExecSpace, Mpi> {
   std::vector<std::function<void()>> comms_;
 
   // wait variables
-  std::vector<MPI_Request> reqs_;
+  std::vector<Req<Mpi>> reqs_;
   std::vector<std::function<void()>> postWaits_;
-  std::vector<std::shared_ptr<ViewHolderBase>> views_;
 };
 
-template <KokkosExecutionSpace ExecSpace>
-void wait(Handle<ExecSpace, Mpi> &handle) {
-  MPI_Waitall(handle.reqs_.size(), handle.reqs_.data(), MPI_STATUSES_IGNORE);
-  handle.reqs_.clear();
-  // std::cerr << __FILE__ << ":" << __LINE__ << " MPI_Waitall done\n";
-  for (const auto &f : handle.postWaits_) {
-    f();
-  }
-  // std::cerr << __FILE__ << ":" << __LINE__ << " postWaits_.clear()...\n";
-  handle.postWaits_.clear();
-  // views_.clear();
-  // std::cerr << __FILE__ << ":" << __LINE__ << " wait() done\n";
-}
-
 }  // namespace KokkosComm
diff --git a/src/mpi/KokkosComm_mpi_plan.hpp b/src/mpi/KokkosComm_mpi_plan.hpp
@@ -16,26 +16,17 @@
 
 #pragma once
 
-namespace KokkosComm {
+namespace KokkosComm::Impl {
 
-template <Dispatch DISPATCH, KokkosExecutionSpace ExecSpace>
-class Plan<DISPATCH, ExecSpace, Mpi> {
- public:
-  using execution_space = ExecSpace;
-  using handle_type     = Handle<execution_space, Mpi>;
-  Plan(const execution_space &space, MPI_Comm comm, DISPATCH d) : handle_(space, comm) {
-    d(handle_);
-    handle_.impl_run();
+template <Dispatch TDispatch, KokkosExecutionSpace ExecSpace>
+struct Plan<TDispatch, ExecSpace, Mpi> {
+  Plan(Handle<ExecSpace, Mpi> &handle, TDispatch d) {
+    d(handle);
+    handle.impl_run();
+    reqs = handle.impl_reqs();
   }
 
-  Plan(const execution_space &space, DISPATCH d) : Plan(space, MPI_COMM_WORLD, d) {}
-  Plan(MPI_Comm comm, DISPATCH d) : Plan(Kokkos::DefaultExecutionSpace(), comm, d) {}
-  Plan(DISPATCH d) : Plan(Kokkos::DefaultExecutionSpace(), MPI_COMM_WORLD, d) {}
-
-  handle_type handle() const { return handle_; }
-
- private:
-  handle_type handle_;
+  std::vector<Req<Mpi>> reqs;
 };
 
-}  // namespace KokkosComm
+}  // namespace KokkosComm::Impl
diff --git a/src/mpi/KokkosComm_mpi_req.hpp b/src/mpi/KokkosComm_mpi_req.hpp
@@ -16,14 +16,17 @@
 
 #pragma once
 
-#include <memory>
 #include <vector>
+#include <utility>
 
-#include "KokkosComm_include_mpi.hpp"
+#include "KokkosComm_fwd.hpp"
+
+#include "KokkosComm_mpi.hpp"
 
 namespace KokkosComm {
 
-class Req {
+template <>
+class Req<Mpi> {
   // a type-erased view. Request uses these to keep temporary views alive for
   // the lifetime of "Immediate" MPI operations
   struct ViewHolderBase {
@@ -38,27 +41,53 @@ class Req {
   struct Record {
     Record() : req_(MPI_REQUEST_NULL) {}
     MPI_Request req_;
-    std::vector<std::shared_ptr<ViewHolderBase>> until_waits_;
+    std::vector<std::function<void()>> postWaits_;
   };
 
  public:
   Req() : record_(std::make_shared<Record>()) {}
 
   MPI_Request &mpi_req() { return record_->req_; }
 
-  void wait() {
-    MPI_Wait(&(record_->req_), MPI_STATUS_IGNORE);
-    record_->until_waits_.clear();  // drop any views we're keeping alive until wait()
-  }
-
   // keep a reference to this view around until wait() is called
   template <typename View>
   void keep_until_wait(const View &v) {
-    record_->until_waits_.push_back(std::make_shared<ViewHolder<View>>(v));
+    record_->postWaits_.push_back([v]() {} /* capture v into a no-op lambda that does nothing*/);
   }
 
  private:
   std::shared_ptr<Record> record_;
+
+  friend void wait(Req<Mpi> &req);
+  friend void wait_all(std::vector<Req<Mpi>> &reqs);
+  friend int wait_any(std::vector<Req<Mpi>> &reqs);
 };
 
+inline void wait(Req<Mpi> &req) {
+  std::cerr << __FILE__ << ":" << __LINE__ << " wait on " << req.mpi_req() << "\n";
+  MPI_Wait(&req.mpi_req(), MPI_STATUS_IGNORE);
+  for (auto &f : req.record_->postWaits_) {
+    f();
+  }
+  req.record_->postWaits_.clear();
+}
+
+inline void wait_all(std::vector<Req<Mpi>> &reqs) {
+  std::cerr << __FILE__ << ":" << __LINE__ << " wait all\n";
+  for (Req<Mpi> &req : reqs) {
+    wait(req);
+  }
+}
+
+inline int wait_any(std::vector<Req<Mpi>> &reqs) {
+  for (size_t i = 0; i < reqs.size(); ++i) {
+    int completed;
+    MPI_Test(&(reqs[i].mpi_req()), &completed, MPI_STATUS_IGNORE);
+    if (completed) {
+      return true;
+    }
+  }
+  return false;
+}
+
 }  // namespace KokkosComm
diff --git a/src/mpi/impl/KokkosComm_mpi_irecv_datatype.hpp b/src/mpi/impl/KokkosComm_mpi_irecv_datatype.hpp
@@ -26,9 +26,9 @@ void irecv_datatype(HandleTy &h, RecvView rv, int src, int tag) {
   h.impl_add_pre_comm_fence();
 
   h.impl_add_comm([&h, rv, src, tag]() {
-    MPI_Request req;
-    MPI_Irecv(KokkosComm::data_handle(rv), 1, view_mpi_type(rv), src, tag, h.mpi_comm(), &req);
-    h.impl_track_mpi_request(req);
+    Req<Mpi> req;
+    MPI_Irecv(KokkosComm::data_handle(rv), 1, view_mpi_type(rv), src, tag, h.mpi_comm(), &req.mpi_req());
+    h.impl_track_req(req);
   });
 }
 
diff --git a/src/mpi/impl/KokkosComm_mpi_isend_datatype.hpp b/src/mpi/impl/KokkosComm_mpi_isend_datatype.hpp
@@ -26,9 +26,9 @@ void isend_datatype(HandleTy &h, SendView sv, int dst, int tag) {
   h.impl_add_pre_comm_fence();
 
   h.impl_add_comm([&h, sv, dst, tag]() {
-    MPI_Request req;
-    MPI_Isend(KokkosComm::data_handle(sv), 1, view_mpi_type(sv), dst, tag, h.mpi_comm(), &req);
-    h.impl_track_mpi_request(req);
+    Req<Mpi> req;
+    MPI_Isend(KokkosComm::data_handle(sv), 1, view_mpi_type(sv), dst, tag, h.mpi_comm(), &req.mpi_req());
+    h.impl_track_req(req);
   });
 }
 
diff --git a/unit_tests/test_isendirecv.cpp b/unit_tests/test_isendirecv.cpp
