Merge pull request #208 from LLNL/multisendrecv

MultiSendRecv operation

Author: ndryden

benchmark/benchmark_ops.cpp

@@ -124,6 +124,18 @@ void run_benchmark(cxxopts::ParseResult& parsed_opts) {
       }
     }
   }
+  if (Op == Al::AlOperation::multisendrecv) {
+    // Send to and receive from every other rank.
+    op_options.srcs = std::vector<int>(comm_wrapper.size() - 1);
+    op_options.dests = std::vector<int>(comm_wrapper.size() - 1);
+    for (int i = 0, rank = 0; rank < comm_wrapper.size(); ++rank) {
+      if (rank != comm_wrapper.rank()) {
+        op_options.srcs[i] = rank;
+        op_options.dests[i] = rank;
+        ++i;
+      }
+    }
+  }
 
   size_t num_iters = parsed_opts["num-iters"].as<size_t>();
   size_t num_warmup = parsed_opts["num-warmup"].as<size_t>();
@@ -135,6 +147,11 @@ void run_benchmark(cxxopts::ParseResult& parsed_opts) {
       op_options.recv_counts = op_options.send_counts;
       op_options.recv_displs = op_options.send_displs;
     }
+    if (Op == Al::AlOperation::multisendrecv) {
+      // Set up to be similar to vector operations.
+      op_options.send_counts = std::vector<size_t>(comm_wrapper.size() - 1, size);
+      op_options.recv_counts = op_options.send_counts;
+    }
 
     OpDispatcher<Backend, T> op_runner(op, op_options);
     size_t in_size = op_runner.get_input_size(size, comm_wrapper.comm());

benchmark/run_benchmarks.py

@@ -86,19 +86,20 @@
 pt2pt_ops = [OpDesc('send', inplace=False, min_procs=2),
              OpDesc('recv', inplace=False, min_procs=2),
              OpDesc('sendrecv', inplace=False, min_procs=2)]
+other_ops = [OpDesc('multisendrecv')]
 
 # Full set of cases.
 test_cases = {
     'mpi': {
-        'ops': coll_ops + vector_coll_ops + pt2pt_ops,
+        'ops': coll_ops + vector_coll_ops + pt2pt_ops + other_ops,
         'datatypes': mpi_datatypes
     },
     'nccl': {
-        'ops': coll_ops + vector_coll_ops + pt2pt_ops,
+        'ops': coll_ops + vector_coll_ops + pt2pt_ops + other_ops,
         'datatypes': nccl_datatypes
     },
     'ht': {
-        'ops': coll_ops + vector_coll_ops + pt2pt_ops,
+        'ops': coll_ops + vector_coll_ops + pt2pt_ops + other_ops,
         'datatypes': mpi_datatypes
     },
 }

docs/comm.rst

@@ -87,10 +87,12 @@ Send and Receive
 
 These are the most basic communication operations, and need to be paired: one process sends data and the other process receives the data.
 
-There are some subtleties here that one should be aware of.
-Unlike MPI, Aluminum does not support tagged messages.
-You should properly order your sends and receives, so that you start your sends and receives in the same order.
-Finally, note that sends may (but are not guaranteed) to complete *before* the data has been sent or before it has been received (it is, however, safe to read or modify a send buffer after the send has completed).
+There are some subtleties here that one should be aware of:
+
+* Unlike MPI, Aluminum does not support tagged messages.
+* Also unlike MPI, a receive will only match a send with the exact same ``count`` (whereas MPI treats the receive count as an upper bound).
+* You should properly order your sends and receives, so that you start your sends and receives in the same order.
+* Finally, note that sends may (but are not guaranteed) to complete *before* the data has been sent or before it has been received (it is, however, safe to read or modify a send buffer after the send has completed).
 
 There are no special in-place versions of send or receive, as they only have one buffer argument.
 
@@ -116,6 +118,20 @@ There is an in-place version of ``SendRecv`` (this is similar to ``MPI_Sendrecv_
 * :cpp:func:`Al::SendRecv()`
 * :cpp:func:`Al::NonblockingSendRecv()`
 
+.. _multisendrecv
+
+MultiSendRecv
+^^^^^^^^^^^^^
+
+This implements an arbitrary sequence of ``Send`` and ``Recv`` operations as a single operation.
+This is similar to an ``Alltoall`` collective, but does not require that all ranks in a communicator participate.
+Likewise, unlike a sequence of ``SendRecv`` operations, this does not require all ranks to both send and receive; i.e., there may be asymmetric communication.
+
+There is an in-place version of ``MultiSendRecv``, which differs slightly from a regular ``MultiSendRecv``: It functions exactly like a sequence of in-place ``SendRecv`` operations.
+
+* :cpp:func:`Al::MultiSendRecv()`
+* :cpp:func:`Al::NonblockingMultiSendRecv()`
+
 Collective Operations
 ---------------------
 

include/Al.hpp

@@ -1404,7 +1404,7 @@ void NonblockingScatterv(
 /**
  * Send a point-to-point message.
  *
- * See \verbatim :embed:rst:inline :ref:`Send and Recv <send-and-recv>`. \endverbatim
+ * See \verbatim embed:rst:inline :ref:`Send and Recv <send-and-recv>`. \endverbatim
  *
  * @param[in] sendbuf Buffer containing the local data to send.
  * @param[in] count Length of \p sendbuf in elements of type `T`.
@@ -1439,7 +1439,7 @@ void NonblockingSend(const T* sendbuf, size_t count, int dest,
 /**
  * Receive a point-to-point message.
  *
- * See \verbatim :embed:rst:inline :ref:`Send and Recv <send-and-recv>`. \endverbatim
+ * See \verbatim embed:rst:inline :ref:`Send and Recv <send-and-recv>`. \endverbatim
  *
  * @param[out] recvbuf Buffer to receive the sent data.
  * @param[in] count Length of \p recvbuf in elements of type `T`.
@@ -1474,7 +1474,7 @@ void NonblockingRecv(T* recvbuf, size_t count, int src,
 /**
  * Perform a simultaneous Send() and Recv().
  *
- * See \verbatim :embed:rst:inline :ref:`SendRecv <sendrecv>`. \endverbatim
+ * See \verbatim embed:rst:inline :ref:`SendRecv <sendrecv>`. \endverbatim
  *
  * @param[in] sendbuf Buffer containing the local data to send.
  * @param[in] send_count Length of \p sendbuf in elements of type `T`.
@@ -1497,7 +1497,7 @@ void SendRecv(const T* sendbuf, size_t send_count, int dest,
 /**
  * Perform an \verbatim embed:rst:inline :ref:`in-place <comm-inplace>` \endverbatim SendRecv().
  *
- * @param[in,out] buffer Input and output buffer initially contaiuning the
+ * @param[in,out] buffer Input and output buffer initially containing the
  * local data to send. Will be replaced with the received data.
  * @param[in] count Length of data to send and receive. \p buffer should
  * be `count` elements.
@@ -1522,8 +1522,8 @@ void SendRecv(T* buffer, size_t count, int dest, int src,
  * @param[out] recvbuf Buffer to receive the sent data.
  * @param[in] recv_count Length of \p recvbuf in elements of type `T`.
  * @param[in] src Rank in comm to receive from.
- * @param[out] req Request object for the asynchronous operation.
  * @param[in] comm Communicator to send/recv within.
+ * @param[out] req Request object for the asynchronous operation.
  */
 template <typename Backend, typename T>
 void NonblockingSendRecv(const T* sendbuf, size_t send_count, int dest,
@@ -1548,8 +1548,8 @@ void NonblockingSendRecv(const T* sendbuf, size_t send_count, int dest,
  * be `count` elements.
  * @param[in] dest Rank in comm to send to.
  * @param[in] src Rank in comm to receive from.
- * @param[out] req Request object for the asynchronous operation.
  * @param[in] comm Communicator to send/recv within.
+ * @param[out] req Request object for the asynchronous operation.
  */
 template <typename Backend, typename T>
 void NonblockingSendRecv(T* buffer, size_t count, int dest, int src,
@@ -1560,6 +1560,116 @@ void NonblockingSendRecv(T* buffer, size_t count, int dest, int src,
   Backend::template NonblockingSendRecv<T>(buffer, count, dest, src, comm, req);
 }
 
+/**
+ * Perform an arbitrary sequence of Send() and Recv() operations.
+ *
+ * See \verbatim embed:rst:inline :ref:`MultiSendRecv <multisendrecv>`. \endverbatim
+ *
+ * @param[in] send_buffers Vector of buffers containing the local data to send.
+ * @param[in] send_counts Vector of the lengths of each buffer in
+ * \p send_buffers in elements of type `T`.
+ * @param[in] dests Vector of the destination rank to send each buffer to.
+ * @param[out] recv_buffers Vector of buffers to receive data in.
+ * @param[in] recv_counts Vector of the lengths of each buffer in
+ * \p recv_buffers in elements of type `T`.
+ * @param[in] srcs Vector of the ranks to receive from.
+ * @param[in] comm Communicator to send/recv within.
+ */
+template <typename Backend, typename T>
+void MultiSendRecv(std::vector<const T*> send_buffers,
+                   std::vector<size_t> send_counts,
+                   std::vector<int> dests,
+                   std::vector<T*> recv_buffers,
+                   std::vector<size_t> recv_counts,
+                   std::vector<int> srcs,
+                   typename Backend::comm_type& comm) {
+  internal::trace::record_op<Backend, T>("multisendrecv", comm,
+                                         send_buffers, send_counts, dests,
+                                         recv_buffers, recv_counts, srcs);
+  Backend::template MultiSendRecv<T>(send_buffers, send_counts, dests,
+                                     recv_buffers, recv_counts, srcs,
+                                     comm);
+}
+
+/**
+ * Perform an \verbatim embed:rst:inline :ref:`in-place <comm-inplace>` \endverbatim MultiSendRecv().
+ *
+ * @param[in, out] buffers Vector of input and output buffers initially
+ * containing the local data to send. Will be replaced with the received
+ * data.
+ * @param[in] counts Vector of the lengths of data to send and receive.
+ * @param[in] dests Vector of the destination rank to send each buffer to.
+ * @param[in] srcs Vector of the ranks to receive from.
+ * @param[in] comm Communicator to send/recv within.
+ */
+template <typename Backend, typename T>
+void MultiSendRecv(std::vector<T*> buffers,
+                   std::vector<size_t> counts,
+                   std::vector<int> dests,
+                   std::vector<int> srcs,
+                   typename Backend::comm_type& comm) {
+  internal::trace::record_op<Backend, T>("multisendrecv", comm,
+                                         buffers, counts, dests, srcs);
+  Backend::template MultiSendRecv<T>(buffers, counts, dests, srcs, comm);
+}
+
+/**
+ * Perform a \verbatim embed:rst:inline :ref:`non-blocking <comm-nonblocking>` \endverbatim MultiSendRecv().
+ *
+ * @param[in] send_buffers Vector of buffers containing the local data to send.
+ * @param[in] send_counts Vector of the lengths of each buffer in
+ * \p send_buffers in elements of type `T`.
+ * @param[in] dests Vector of the destination rank to send each buffer to.
+ * @param[out] recv_buffers Vector of buffers to receive data in.
+ * @param[in] recv_counts Vector of the lengths of each buffer in
+ * \p recv_buffers in elements of type `T`.
+ * @param[in] srcs Vector of the ranks to receive from.
+ * @param[in] comm Communicator to send/recv within.
+ * @param[out] req Request object for the asynchronous operation.
+ */
+template <typename Backend, typename T>
+void NonblockingMultiSendRecv(std::vector<const T*> send_buffers,
+                              std::vector<size_t> send_counts,
+                              std::vector<int> dests,
+                              std::vector<T*> recv_buffers,
+                              std::vector<size_t> recv_counts,
+                              std::vector<int> srcs,
+                              typename Backend::comm_type& comm,
+                              typename Backend::req_type& req) {
+  internal::trace::record_op<Backend, T>("nonblocking-multisendrecv", comm,
+                                         send_buffers, send_counts, dests,
+                                         recv_buffers, recv_counts, srcs);
+  Backend::template NonblockingMultiSendRecv<T>(send_buffers, send_counts, dests,
+                                                recv_buffers, recv_counts, srcs,
+                                                comm, req);
+}
+
+/**
+ * Perform a \verbatim embed:rst:inline :ref:`non-blocking <comm-nonblocking>` \endverbatim
+ * \verbatim embed:rst:inline :ref:`in-place <comm-inplace>` \endverbatim MultiSendRecv().
+ *
+ * @param[in, out] buffers Vector of input and output buffers initially
+ * containing the local data to send. Will be replaced with the received
+ * data.
+ * @param[in] counts Vector of the lengths of data to send and receive.
+ * @param[in] dests Vector of the destination rank to send each buffer to.
+ * @param[in] srcs Vector of the ranks to receive from.
+ * @param[in] comm Communicator to send/recv within.
+ * @param[out] req Request object for the asynchronous operation.
+ */
+template <typename Backend, typename T>
+void NonblockingMultiSendRecv(std::vector<T*> buffers,
+                              std::vector<size_t> counts,
+                              std::vector<int> dests,
+                              std::vector<int> srcs,
+                              typename Backend::comm_type& comm,
+                              typename Backend::req_type& req) {
+  internal::trace::record_op<Backend, T>("nonblocking-multisendrecv", comm,
+                                         buffers, counts, dests, srcs);
+  Backend::template NonblockingMultiSendRecv<T>(buffers, counts, dests, srcs,
+                                                comm, req);
+}
+
 /**
  * Return true if the asynchronous operation associated with \p req
  * has completed.

include/aluminum/ht/CMakeLists.txt

@@ -10,6 +10,7 @@ set_source_path(THIS_DIR_HEADERS
   communicator.hpp
   gather.hpp
   gatherv.hpp
+  multisendrecv.hpp
   reduce.hpp
   reduce_scatter.hpp
   reduce_scatterv.hpp

include/aluminum/ht/communicator.hpp

@@ -37,6 +37,9 @@ namespace Al {
 namespace internal {
 namespace ht {
 
+// Define the tag that point-to-point operations will use.
+constexpr int pt2pt_tag = 2;
+
 /** Communicator for host-transfer operations. */
 class HostTransferCommunicator: public MPICommAndStreamWrapper<AlGpuStream_t> {
  public: