Add send/receive payload tracking for non-shuffle collectives

cpp/include/rapidsmpf/shuffler/shuffler.hpp

@@ -344,8 +344,6 @@ class Shuffler {
     std::vector<detail::ChunkID> outbound_chunk_counter_;  ///< indexed by Rank
     std::atomic<detail::ChunkID> chunk_id_counter_{0};
 
-    std::shared_ptr<Statistics> statistics_;
-
     // For notifications.
     mutable std::mutex mutex_;
     std::condition_variable cv_;

cpp/src/coll/allgather.cpp

@@ -163,6 +163,7 @@ ProgressThread::ProgressState AllGather::event_loop() {
      */
     Rank const dst = (comm_->rank() + 1) % comm_->nranks();
     Rank const src = (comm_->rank() + comm_->nranks() - 1) % comm_->nranks();
+    auto const& statistics = comm_->progress_thread()->statistics();
     // GPU data sends and metadata sends can be arbitrarily interleaved. To allow reuse of
     // `op_id` once `wait_and_extract()` returns, we rely on a number of invariants
     // enforced by the communication scheme.
@@ -223,6 +224,11 @@ ProgressThread::ProgressState AllGather::event_loop() {
                 // of insertions from that rank.
                 mark_finish(chunk->sequence());
             } else {
+                if (chunk->data_size() > 0) {
+                    statistics->add_bytes_stat(
+                        "allgather-payload-send", chunk->data_size()
+                    );
+                }
                 auto buf = chunk->release_data_buffer();
                 sent_posted_.emplace_back(std::move(chunk));
                 sent_futures_.emplace_back(
@@ -261,6 +267,9 @@ ProgressThread::ProgressState AllGather::event_loop() {
             if (!chunk->is_ready()) {
                 break;
             }
+            if (chunk->data_size() > 0) {
+                statistics->add_bytes_stat("allgather-payload-recv", chunk->data_size());
+            }
             auto buf = chunk->release_data_buffer();
             receive_posted_.emplace_back(std::move(chunk));
             receive_futures_.emplace_back(comm_->recv(src, gpu_data_tag, std::move(buf)));

cpp/src/coll/allreduce.cpp

@@ -59,8 +59,12 @@ AllReduce::AllReduce(
     out_buffer_->rebind_stream(in_buffer_->stream());
     // Note: after this copy, we must check out_buffer's write event before receiving into
     // in_buffer. See StartPreRemainder in the event loop.
-    // TODO: make Communicator statistics-aware, and pass its statistics instance here.
-    buffer_copy(Statistics::disabled(), *out_buffer_, *in_buffer_, in_buffer_->size);
+    buffer_copy(
+        comm_->progress_thread()->statistics(),
+        *out_buffer_,
+        *in_buffer_,
+        in_buffer_->size
+    );
 
     auto const rank = comm_->rank();
     if (rank < 2 * non_pow2_remainder_) {
@@ -134,6 +138,7 @@ bool AllReduce::is_ready() const noexcept {
 ProgressThread::ProgressState AllReduce::event_loop() {
     Rank const rank = comm_->rank();
     bool const is_even = rank % 2 == 0;
+    auto const& statistics = comm_->progress_thread()->statistics();
     // We only need a single stage ID because of no-message-overtaking guarantees in the
     // communicator. We could use multiple stage IDs for each round of the exchange, but
     // that would break once we have more than 256 participating ranks: there are only
@@ -173,6 +178,7 @@ ProgressThread::ProgressState AllReduce::event_loop() {
                 if (!out_buffer_->is_latest_write_done()) {
                     break;
                 }
+                statistics->add_bytes_stat("allreduce-payload-send", out_buffer_->size);
                 send_future_ = comm_->send(std::move(out_buffer_), rank + 1, tag);
             } else {
                 // The constructor copies in_buffer_ to out_buffer_ on in_buffer's
@@ -186,6 +192,7 @@ ProgressThread::ProgressState AllReduce::event_loop() {
                 {
                     break;
                 }
+                statistics->add_bytes_stat("allreduce-payload-recv", in_buffer_->size);
                 recv_future_ = comm_->recv(rank - 1, tag, std::move(in_buffer_));
             }
             phase_.store(Phase::CompletePreRemainder, std::memory_order_release);
@@ -237,7 +244,9 @@ ProgressThread::ProgressState AllReduce::event_loop() {
             {
                 break;
             }
+            statistics->add_bytes_stat("allreduce-payload-recv", in_buffer_->size);
             recv_future_ = comm_->recv(stage_partner_, tag, std::move(in_buffer_));
+            statistics->add_bytes_stat("allreduce-payload-send", out_buffer_->size);
             send_future_ = comm_->send(std::move(out_buffer_), stage_partner_, tag);
             phase_.store(Phase::CompleteButterfly, std::memory_order_release);
             break;
@@ -269,11 +278,13 @@ ProgressThread::ProgressState AllReduce::event_loop() {
                 if (!out_buffer_->is_latest_write_done()) {
                     break;
                 }
+                statistics->add_bytes_stat("allreduce-payload-recv", out_buffer_->size);
                 recv_future_ = comm_->recv(rank + 1, tag, std::move(out_buffer_));
             } else {
                 if (!out_buffer_->is_latest_write_done()) {
                     break;
                 }
+                statistics->add_bytes_stat("allreduce-payload-send", out_buffer_->size);
                 send_future_ = comm_->send(std::move(out_buffer_), rank - 1, tag);
             }
             phase_.store(Phase::CompletePostRemainder, std::memory_order_release);

cpp/src/coll/sparse_alltoall.cpp

@@ -157,10 +157,12 @@ std::vector<PackedData> SparseAlltoall::extract(Rank src) {
 void SparseAlltoall::send_ready_messages() {
     Tag const metadata_tag{op_id_, 0};
     Tag const payload_tag{op_id_, 1};
+    auto const& statistics = comm_->progress_thread()->statistics();
     for (auto& chunk : outgoing_.extract_ready()) {
         auto const dst = chunk->destination();
         fire_and_forget_.push_back(comm_->send(chunk->serialize(), dst, metadata_tag));
         if (chunk->data_size() > 0) {
+            statistics->add_bytes_stat("sparsealltoall-payload-send", chunk->data_size());
             fire_and_forget_.push_back(
                 comm_->send(chunk->release_data_buffer(), dst, payload_tag)
             );
@@ -197,6 +199,7 @@ void SparseAlltoall::receive_metadata_messages() {
 
 void SparseAlltoall::receive_data_messages() {
     Tag const payload_tag{op_id_, 1};
+    auto const& statistics = comm_->progress_thread()->statistics();
     for (auto& [src, state] : source_states_) {
         std::ptrdiff_t processed = 0;
         auto& queue = state.incoming;
@@ -208,6 +211,9 @@ void SparseAlltoall::receive_data_messages() {
             if (chunk->data_size() == 0) {
                 state.chunks.push_back(std::move(chunk));
             } else {
+                statistics->add_bytes_stat(
+                    "sparsealltoall-payload-recv", chunk->data_size()
+                );
                 auto buffer = chunk->release_data_buffer();
                 receive_posted_.push_back(std::move(chunk));
                 receive_futures_.push_back(

cpp/src/shuffler/shuffler.cpp

@@ -89,7 +89,7 @@ class Shuffler::Progress {
         RAPIDSMPF_NVTX_SCOPED_RANGE_VERBOSE("Shuffler.Progress", p_iters++);
         auto const t0_event_loop = Clock::now();
 
-        auto& stats = *shuffler_.statistics_;
+        auto const& stats = shuffler_.comm_->progress_thread()->statistics();
 
         // Submit outgoing chunks to the metadata payload exchange
         {
@@ -114,7 +114,7 @@ class Shuffler::Progress {
 
                 for (auto const& chunk : ready_chunks) {
                     if (chunk.data_size() > 0) {
-                        stats.add_bytes_stat("shuffle-payload-send", chunk.data_size());
+                        stats->add_bytes_stat("shuffle-payload-send", chunk.data_size());
                     }
                 }
 
@@ -123,7 +123,7 @@ class Shuffler::Progress {
 
                 shuffler_.mpe_->send(std::move(messages));
             }
-            stats.add_duration_stat(
+            stats->add_duration_stat(
                 "event-loop-submit-outgoing", Clock::now() - t0_submit_outgoing
             );
         }
@@ -149,18 +149,18 @@ class Shuffler::Progress {
                 );
 
                 if (chunk.data_size() > 0) {
-                    stats.add_bytes_stat("shuffle-payload-recv", chunk.data_size());
+                    stats->add_bytes_stat("shuffle-payload-recv", chunk.data_size());
                 }
 
                 shuffler_.insert_into_received(std::move(chunk));
             }
 
-            stats.add_duration_stat(
+            stats->add_duration_stat(
                 "event-loop-process-communication", Clock::now() - t0_process_comm
             );
         }
 
-        stats.add_duration_stat("event-loop-total", Clock::now() - t0_event_loop);
+        stats->add_duration_stat("event-loop-total", Clock::now() - t0_event_loop);
 
         // Signal the MPE that no more messages will be sent once all application
         // messages have been flushed from to_send_ into the MPE.
@@ -249,13 +249,12 @@ Shuffler::Shuffler(
                             br_->reserve_or_fail(size, MEMORY_TYPES)
                         );
                     },
-                    br_->statistics()
+                    comm_->progress_thread()->statistics()
                 )
       },
       local_partitions_{local_partitions(comm_, total_num_partitions, partition_owner)},
       finish_counter_{comm_->nranks(), safe_cast<PartID>(local_partitions_.size())},
       outbound_chunk_counter_(safe_cast<std::size_t>(comm_->nranks()), 0),
-      statistics_{br_->statistics()},
       finished_callback_{std::move(finished_callback)} {
     RAPIDSMPF_EXPECTS(
         total_num_partitions > 0, "number of partitions must be strictly positive"
@@ -328,10 +327,6 @@ void Shuffler::insert(detail::Chunk&& chunk) {
     }
     if (dst_rank == comm_->rank()) {
         // this is a local chunk, so we can move it straight to received.
-        if (chunk.is_data_buffer_set()) {
-            statistics_->add_bytes_stat("shuffle-payload-send", chunk.data_size());
-            statistics_->add_bytes_stat("shuffle-payload-recv", chunk.data_size());
-        }
         insert_into_received(std::move(chunk));
     } else {
         // this is a remote chunk, so we need to send it

cpp/tests/environment.hpp

@@ -4,10 +4,45 @@
  */
 #pragma once
 
+#include <utility>
+
 #include <gtest/gtest.h>
 #include <mpi.h>
 
 #include <rapidsmpf/communicator/communicator.hpp>
+#include <rapidsmpf/statistics.hpp>
+
+/// @brief Statistics wrapper that clears statistics counters on construction and
+/// destruction
+///
+/// Use this if you want to track statistics for a single test from the shared
+/// communicator stats object.
+class ClearedStatistics {
+  public:
+    explicit ClearedStatistics(std::shared_ptr<rapidsmpf::Statistics> statistics)
+        : statistics_{std::move(statistics)}, was_enabled_{statistics_->enabled()} {
+        statistics_->clear();
+        statistics_->enable();
+    }
+
+    ~ClearedStatistics() {
+        statistics_->clear();
+        if (!was_enabled_) {
+            statistics_->disable();
+        }
+    }
+
+    ClearedStatistics(ClearedStatistics const&) = delete;
+    ClearedStatistics& operator=(ClearedStatistics const&) = delete;
+
+    [[nodiscard]] rapidsmpf::Statistics* operator->() const noexcept {
+        return statistics_.get();
+    }
+
+  private:
+    std::shared_ptr<rapidsmpf::Statistics> statistics_;
+    bool was_enabled_;
+};
 
 enum class TestEnvironmentType : int {
     MPI,

cpp/tests/test_allgather.cpp

@@ -160,6 +160,42 @@ TEST_P(AllGatherTest, basic_allgather) {
     }
 }
 
+TEST_F(BaseAllGatherTest, payload_statistics) {
+    auto const& comm = GlobalEnvironment->comm_;
+    ClearedStatistics statistics{comm->progress_thread()->statistics()};
+    constexpr int n_elements = 7;
+    constexpr int n_inserts = 3;
+
+    AllGather allgather{comm, 0, br.get()};
+    for (int i = 0; i < n_inserts; ++i) {
+        allgather.insert(
+            i, generate_packed_data(n_elements, gen_offset(i, comm->rank()), stream, *br)
+        );
+    }
+    allgather.insert_finished();
+    auto results =
+        allgather.wait_and_extract(AllGather::Ordered::NO, std::chrono::seconds{30});
+    EXPECT_EQ(results.size(), static_cast<std::size_t>(n_inserts * comm->nranks()));
+
+    auto const expected_count =
+        static_cast<std::size_t>(n_inserts * (comm->nranks() - 1));
+    if (expected_count == 0) {
+        EXPECT_THROW(statistics->get_stat("allgather-payload-send"), std::out_of_range);
+        EXPECT_THROW(statistics->get_stat("allgather-payload-recv"), std::out_of_range);
+    } else {
+        auto const expected_message_size = n_elements * sizeof(int);
+        auto const expected_bytes = expected_count * expected_message_size;
+        auto const send = statistics->get_stat("allgather-payload-send");
+        auto const recv = statistics->get_stat("allgather-payload-recv");
+        EXPECT_EQ(send.count(), expected_count);
+        EXPECT_EQ(send.value(), expected_bytes);
+        EXPECT_EQ(send.max(), expected_message_size);
+        EXPECT_EQ(recv.count(), expected_count);
+        EXPECT_EQ(recv.value(), expected_bytes);
+        EXPECT_EQ(recv.max(), expected_message_size);
+    }
+}
+
 class AllGatherOrderedTest : public BaseAllGatherTest,
                              public ::testing::WithParamInterface<AllGather::Ordered> {};
 

cpp/tests/test_allreduce.cu

@@ -28,6 +28,7 @@
 #include <rapidsmpf/memory/buffer_resource.hpp>
 #include <rapidsmpf/memory/cuda_memcpy_async.hpp>
 #include <rapidsmpf/memory/memory_type.hpp>
+#include <rapidsmpf/statistics.hpp>
 
 #include "environment.hpp"
 
@@ -363,6 +364,57 @@ TEST_P(AllReduceIntSumTest, basic_allreduce_sum_int) {
     EXPECT_TRUE(allreduce.finished());
 }
 
+TEST_F(BaseAllReduceTest, payload_statistics) {
+    ClearedStatistics statistics{comm->progress_thread()->statistics()};
+    constexpr int n_elements = 11;
+    std::vector<int> data(n_elements, comm->rank());
+    auto in_buffer =
+        make_buffer<int>(br.get(), data.data(), data.size(), MemoryType::HOST);
+    auto const message_size = in_buffer->size;
+    auto reservation = br->reserve_or_fail(message_size, MemoryType::HOST);
+    auto out_buffer = br->make_buffer(message_size, in_buffer->stream(), reservation);
+
+    AllReduce allreduce(
+        GlobalEnvironment->comm_,
+        std::move(in_buffer),
+        std::move(out_buffer),
+        OpID{0},
+        rapidsmpf::coll::detail::make_host_reduce_operator<int>(SumOp<int>{})
+    );
+    std::ignore = allreduce.wait_and_extract();
+
+    // For checking the statistics we need to know the number of rounds each rank
+    // participates in. Note: this has to match the communication pattern in the allreduce
+    // implementation.
+    auto nearest_power_of_two = 1;
+    auto butterfly_stages = 0;
+    while (nearest_power_of_two * 2 <= comm->nranks()) {
+        nearest_power_of_two *= 2;
+        ++butterfly_stages;
+    }
+    auto const remainder = comm->nranks() - nearest_power_of_two;
+    auto const rank = comm->rank();
+    auto const expected_count = static_cast<std::size_t>(
+        rank < 2 * remainder ? (rank % 2 == 0 ? 1 : butterfly_stages + 1)
+                             : butterfly_stages
+    );
+
+    if (expected_count == 0) {
+        EXPECT_THROW(statistics->get_stat("allreduce-payload-send"), std::out_of_range);
+        EXPECT_THROW(statistics->get_stat("allreduce-payload-recv"), std::out_of_range);
+    } else {
+        auto const expected_bytes = expected_count * message_size;
+        auto const send = statistics->get_stat("allreduce-payload-send");
+        auto const recv = statistics->get_stat("allreduce-payload-recv");
+        EXPECT_EQ(send.count(), expected_count);
+        EXPECT_EQ(send.value(), expected_bytes);
+        EXPECT_EQ(send.max(), message_size);
+        EXPECT_EQ(recv.count(), expected_count);
+        EXPECT_EQ(recv.value(), expected_bytes);
+        EXPECT_EQ(recv.max(), message_size);
+    }
+}
+
 template <typename T, typename Op, MemoryType MemType>
 struct AllReduceCase {
     using value_type = T;