non-contiguous Impl::irecv

docs/api/core.rst

@@ -36,10 +36,34 @@ Core
 Point-to-point
 --------------
 
+.. cpp:function:: template <KokkosExecutionSpace ExecSpace, KokkosView RecvView> \
+                  Req KokkosComm::irecv(const ExecSpace &space, const RecvView &rv, int src, int tag, MPI_Comm comm)
+
+    Wrapper for ``MPI_Irecv``.
+    The communication operation will be inserted into ``space``.
+    The caller may safely call this function on data previously produced by operations in ``space`` without first fencing ```space```.
+
+    .. warning::
+        Even if ``space`` is fenced after the call to this function, the communication operation is not complete until the ``wait`` operation on the returned ``Req`` is called.
+
+    :param space: The execution space to operate in
+    :param rv: The view to recveive data into
+    :param src: the source rank
+    :param tag: the MPI tag
+    :param comm: the MPI communicator
+    :tparam RecvView: The type of the received-to Kokkos::View
+    :tparam ExecSpace: A Kokkos execution space to operate in
+    :returns: A KokkosComm::Req representing the asynchronous communication and any lifetime-extended views, associated with the provided execution space instance.
+
 .. cpp:function:: template <KokkosComm::CommMode SendMode, KokkosExecutionSpace ExecSpace, KokkosView SendView> \
                   Req KokkosComm::isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm)
 
     Wrapper for ``MPI_Isend``, ``MPI_Irsend`` and ``MPI_Issend``.
+    The communication operation will be inserted into ``space``.
+    The caller may safely call this function on data previously produced by operations in ``space`` without first fencing ```space```.
+
+    .. warning::
+        Even if ``space`` is fenced after the call to this function, the communication operation is not complete until the ``wait`` operation on the returned ``Req`` is called.
 
     :param space: The execution space to operate in
     :param sv: The data to send
@@ -49,7 +73,7 @@ Point-to-point
     :tparam SendMode: A CommMode_ to use. If unspecified, defaults to a synchronous ``MPI_Issend`` if ``KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE`` is defined, otherwise defaults to a standard ``MPI_Isend``.
     :tparam SendView: A Kokkos::View to send
     :tparam ExecSpace: A Kokkos execution space to operate in
-    :returns: A KokkosComm::Req representing the asynchronous communication and any lifetime-extended views.
+    :returns: A KokkosComm::Req representing the asynchronous communication and any lifetime-extended views, associated with the provided execution space instance.
 
 .. cpp:function:: template <KokkosComm::CommMode SendMode, KokkosExecutionSpace ExecSpace, KokkosView SendView> \
                   void KokkosComm::send(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm)
@@ -125,18 +149,27 @@ Related Types
 
 .. cpp:class:: KokkosComm::Req
 
-    A wrapper around an MPI_Request that can also extend the lifetime of Views.
+    A communication handle representing an asychronous communication and an associated Kokkos execution space instance. The handle is scoped to the space instance used in the communication call that produced it. 
+
+
 
     .. cpp:function:: MPI_Request &KokkosComm::Req::mpi_req()
 
         Retrieve a reference to the held MPI_Request.
 
     .. cpp:function:: void KokkosComm::Req::wait()
 
-        Call MPI_Wait on the held MPI_Request and drop copies of any previous arguments to Req::keep_until_wait().
+        Require that the communication be completed before any further work can be exected in the associated execution space instance. May or may not fence. Consider the following example.
+
+        .. code-block:: c++
+          :linenos:
 
-    .. cpp:function:: template<typename View> \
-                      void KokkosComm::Req::keep_until_wait(const View &v)
+          using KC = KokkosComm;
+          Kokkos::parallel_for(space, ...);
+          auto req = KC::isend(space, ...); // isend 1
+          Kokkos::parallel_for(space, ...); // runs concurrently with isend 1, does not touch send view
+          req.wait();                       // blocks space until isend 1 is complete. May or may not fence.
+          Kokkos::parallel_for(space, ...); // safe to overwrite the send buffer
+          space.fence(); // wait for all to complete
 
-        Extend the lifetime of v at least until Req::wait() is called.
-        This is useful to prevent a View from being destroyed during an asynchronous MPI operation.
+        Here, ``parallel_for`` on line 6 can overwrite the send buffer because ``req.wait()`` means that isend 1 must be done before additional work can be done in ``space``. This MAY be achieved by an internal call to ``space.fence()``, but some other mechanism may be used. If the host thread wants to be sure that the communication is done, it must separately call ``space.fence()``.

docs/design.rst

@@ -1,15 +1,24 @@
 Design
 ======
 
-Asynchronous MPI operations and view lifetimes
-----------------------------------------------
+Asynchronous MPI operations, View Lifetimes, and Packing
+--------------------------------------------------------
 
-"Immediate" functions (e.g. `isend`) return a `KokkosComm::Req`, which can be `wait()`-ed to block until the input view can be reused. `Req` also manages the lifetimes of any intermediate views needed for packing the data, releasing those views when `wait()` is complete.
+Asynchronous MPI operations use an ``MPI_Request``, which is a handle that can be used to refer to the operation later (e.g., ``MPI_Wait``).
+
+KokkosComm has an analogous concept, `KokkosComm::Req`.
+"Immediate" functions (e.g. `isend`) return a `KokkosComm::Req`, which can be `wait()`-ed to block until the input view can be reused.
+
+There are three consequences
+
+* First, to ensure compatibility with MPI semantics, KokkosComm immediate functions will call the corresponding MPI function before they return.
+* Second, the KokkosComm packing strategy may require that an intermediate view be allocated, and this view needs to have a lifetime at least as long as the communication. The ``KokkosComm::Req::keep_until_wait`` interface allows the `KokkosComm::Req` to hold those views until ``wait`` is called.
+* Third, for asynchronous receive operations, the packing strategy may require that the buffer provided by the underlying MPI operation be further unpacked. The ``KokkosComm::Req::call_and_drop_at_wait`` allows the `KokkosComm::Req` to execute (and then drop) callback functors when ``wait`` is called. For example, `KokkosComm::irecv` uses this functionality to attach an unpacking operation to the `KokkosComm::Req::wait` call.
 
 Non-contiguous Data
 -------------------
 
-- Packer::DeepCopy uses `Kokkos::deep_copy` to handle packing and unpacking of non-contiguous `Kokkos::View`. This requires an intermediate allocation, which only works for Kokkos Views, see `Device Data`_.
+* Packer::DeepCopy uses `Kokkos::deep_copy` to handle packing and unpacking of non-contiguous `Kokkos::View`. This requires an intermediate allocation.
 
 Device Data
 -----------

src/KokkosComm.hpp

@@ -16,8 +16,12 @@
 
 #pragma once
 
-#include "KokkosComm_collective.hpp"
 #include "KokkosComm_version.hpp"
+#include "KokkosComm_concepts.hpp"
+#include "KokkosComm_request.hpp"
+
+#include "KokkosComm_collective.hpp"
+#include "KokkosComm_irecv.hpp"
 #include "KokkosComm_isend.hpp"
 #include "KokkosComm_recv.hpp"
 #include "KokkosComm_send.hpp"
@@ -30,7 +34,7 @@ namespace KokkosComm {
 
 template <CommMode SendMode = CommMode::Default, KokkosExecutionSpace ExecSpace, KokkosView SendView>
 Req isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm) {
-  return Impl::isend<SendMode>(space, sv, dest, tag, comm);
+  return Req(Impl::isend<SendMode>(space, sv, dest, tag, comm));
 }
 
 template <CommMode SendMode = CommMode::Default, KokkosExecutionSpace ExecSpace, KokkosView SendView>
@@ -43,4 +47,9 @@ void recv(const ExecSpace &space, RecvView &sv, int src, int tag, MPI_Comm comm)
   return Impl::recv(space, sv, src, tag, comm);
 }
 
+template <KokkosExecutionSpace ExecSpace, KokkosView RecvView>
+Req irecv(const ExecSpace &space, RecvView &rv, int src, int tag, MPI_Comm comm) {
+  return Req(Impl::irecv(space, rv, src, tag, comm));
+}
+
 }  // namespace KokkosComm

src/KokkosComm_request.hpp

@@ -18,46 +18,22 @@
 
 #include <memory>
 
-#include "KokkosComm_include_mpi.hpp"
+#include "KokkosComm_request_impl.hpp"
 
 namespace KokkosComm {
 
 class Req {
-  // 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_;
-  };
-
-  struct Record {
-    Record() : req_(MPI_REQUEST_NULL) {}
-    MPI_Request req_;
-    std::vector<std::shared_ptr<ViewHolderBase>> until_waits_;
-  };
-
  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()
-  }
+  Req() : impl_(std::make_shared<Impl::Req>()) {}
+  Req(const std::shared_ptr<Impl::Req> &impl) : impl_(impl) {}
+  MPI_Request &mpi_req() { return impl_->mpi_req(); }
 
-  // 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));
-  }
+  // The communication must be done before the associated Kokkos execution space
+  // can do any further work. The host may or may not be fenced.
+  void wait() { impl_->wait(); }
 
  private:
-  std::shared_ptr<Record> record_;
+  std::shared_ptr<Impl::Req> impl_;
 };
 
-}  // namespace KokkosComm
+}  // namespace KokkosComm

src/impl/KokkosComm_concepts.hpp

@@ -25,5 +25,7 @@ concept KokkosView = Kokkos::is_view_v<T>;
 
 template <typename T>
 concept KokkosExecutionSpace = Kokkos::is_execution_space_v<T>;
+template <typename Fn>
+concept Invokable = std::is_invocable_v<Fn>;
 
 }  // namespace KokkosComm

src/impl/KokkosComm_irecv.hpp

@@ -0,0 +1,75 @@
+//@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 <memory>
+
+#include <Kokkos_Core.hpp>
+
+#include "KokkosComm_pack_traits.hpp"
+#include "KokkosComm_traits.hpp"
+
+// impl
+#include "KokkosComm_include_mpi.hpp"
+
+namespace KokkosComm::Impl {
+
+template <KokkosExecutionSpace ExecSpace, KokkosView RecvView, typename MpiArgs>
+struct IrecvUnpacker {
+  IrecvUnpacker(const ExecSpace &space, RecvView &rv, MpiArgs &args) : space_(space), rv_(rv), args_(args) {}
+
+  void operator()() {
+    Kokkos::Tools::pushRegion("KokkosComm::Impl::IrecvUnpacker::operator()");
+    MpiArgs::packer_type::unpack_into(space_, rv_, args_.view);
+    Kokkos::Tools::popRegion();
+  }
+
+  ExecSpace space_;
+  RecvView rv_;
+  MpiArgs args_;
+};
+
+template <KokkosExecutionSpace ExecSpace, KokkosView RecvView>
+std::shared_ptr<Req> irecv(const ExecSpace &space, RecvView &rv, int src, int tag, MPI_Comm comm) {
+  Kokkos::Tools::pushRegion("KokkosComm::Impl::irecv");
+
+  using KCT  = KokkosComm::Traits<RecvView>;
+  using KCPT = KokkosComm::PackTraits<RecvView>;
+
+  auto req = std::make_shared<Req>();
+
+  if (KCPT::needs_unpack(rv)) {
+    using Packer = typename KCPT::packer_type;
+    using Args   = typename Packer::args_type;
+
+    Args args = Packer::allocate_packed_for(space, "packed", rv);
+    space.fence();  // make sure allocation is done
+    MPI_Irecv(KCT::data_handle(args.view), args.count, args.datatype, src, tag, comm, &req->mpi_req());
+    req->call_and_drop_at_wait(IrecvUnpacker{space, rv, args});
+    // req.wait() promises that communication is done before any future work put into space after wait can proceed.
+    // For MPI, wait uses MPI_Wait to make sure the communication is done, which blocks the host, thereby preventing
+    // later work from being submitted to space, so no space fence is needed at wait.
+  } else {
+    using RecvScalar = typename RecvView::value_type;
+    MPI_Irecv(KCT::data_handle(rv), KCT::span(rv), mpi_type_v<RecvScalar>, src, tag, comm, &req->mpi_req());
+    req->keep_until_wait(rv);
+  }
+  return req;
+
+  Kokkos::Tools::popRegion();
+}
+}  // namespace KokkosComm::Impl

src/impl/KokkosComm_isend.hpp

@@ -16,7 +16,7 @@
 
 #pragma once
 
-#include <iostream>
+#include <memory>
 
 #include <Kokkos_Core.hpp>
 
@@ -35,7 +35,7 @@ template <CommMode SendMode = CommMode::Default, KokkosExecutionSpace ExecSpace,
 KokkosComm::Req isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm) {
   Kokkos::Tools::pushRegion("KokkosComm::Impl::isend");
 
-  KokkosComm::Req req;
+  auto req = std::make_shared<Req>();
 
   using KCT  = KokkosComm::Traits<SendView>;
   using KCPT = KokkosComm::PackTraits<SendView>;
@@ -63,13 +63,13 @@ KokkosComm::Req isend(const ExecSpace &space, const SendView &sv, int dest, int
 
     MpiArgs args = Packer::pack(space, sv);
     space.fence();
-    mpi_isend_fn(KCT::data_handle(args.view), args.count, args.datatype, dest, tag, comm, &req.mpi_req());
-    req.keep_until_wait(args.view);
+    mpi_isend_fn(KCT::data_handle(args.view), args.count, args.datatype, dest, tag, comm, &req->mpi_req());
+    req->keep_until_wait(args.view);
   } else {
     using SendScalar = typename SendView::value_type;
-    mpi_isend_fn(KCT::data_handle(sv), KCT::span(sv), mpi_type_v<SendScalar>, dest, tag, comm, &req.mpi_req());
+    mpi_isend_fn(KCT::data_handle(sv), KCT::span(sv), mpi_type_v<SendScalar>, dest, tag, comm, &req->mpi_req());
     if (KCT::is_reference_counted()) {
-      req.keep_until_wait(sv);
+      req->keep_until_wait(sv);
     }
   }
 

src/impl/KokkosComm_packer.hpp

@@ -23,8 +23,10 @@
 namespace KokkosComm::Impl {
 namespace Packer {
 
-template <KokkosView View>
+template <KokkosView View, typename Packer>
 struct MpiArgs {
+  using packer_type = Packer;  // the type of the packer that produced these arguments
+
   View view;
   MPI_Datatype datatype;
   int count;
@@ -37,7 +39,8 @@ template <KokkosView View>
 struct DeepCopy {
   using non_const_packed_view_type =
       Kokkos::View<typename View::non_const_data_type, Kokkos::LayoutRight, typename View::memory_space>;
-  using args_type = MpiArgs<non_const_packed_view_type>;
+
+  using args_type = MpiArgs<non_const_packed_view_type, DeepCopy<View>>;
 
   template <KokkosExecutionSpace ExecSpace>
   static args_type allocate_packed_for(const ExecSpace &space, const std::string &label, const View &src) {
@@ -71,7 +74,7 @@ struct DeepCopy {
 template <KokkosView View>
 struct MpiDatatype {
   using non_const_packed_view_type = View;
-  using args_type                  = MpiArgs<non_const_packed_view_type>;
+  using args_type                  = MpiArgs<non_const_packed_view_type, MpiDatatype<View>>;
 
   // don't actually allocate - return the provided view, but with
   // a datatype that describes the data in the view