port old code

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Author: nirandaperera

cpp/include/rapidsmpf/statistics.hpp

@@ -537,6 +537,22 @@ class Statistics : public std::enable_shared_from_this<Statistics> {
         MemoryType mem_type, std::size_t nbytes, StreamOrderedTiming&& timing
     );
 
+    /**
+     * @brief Record byte count for a send operation.
+     *
+     * @param src Source memory type.
+     * @param nbytes Number of bytes sent.
+     */
+    void record_send(MemoryType src, std::size_t nbytes);
+
+    /**
+     * @brief Record byte count for a receive operation.
+     *
+     * @param dst Destination memory type.
+     * @param nbytes Number of bytes received.
+     */
+    void record_recv(MemoryType dst, std::size_t nbytes);
+
     /**
      * @brief Get the names of all statistics.
      *

cpp/src/communicator/mpi.cpp

@@ -127,6 +127,7 @@ std::unique_ptr<Communicator::Future> MPI::send(
     RAPIDSMPF_MPI(MPI_Isend(
         msg->data(), safe_cast<int>(msg->size()), MPI_UINT8_T, rank, tag, comm_, &req
     ));
+    statistics_->record_send(MemoryType::HOST, msg->size());
     return std::make_unique<Future>(req, std::move(msg));
 }
 
@@ -139,6 +140,7 @@ std::unique_ptr<Communicator::Future> MPI::send(
     RAPIDSMPF_MPI(MPI_Isend(
         msg->data(), safe_cast<int>(msg->size), MPI_UINT8_T, rank, tag, comm_, &req
     ));
+    statistics_->record_send(msg->mem_type(), msg->size);
     return std::make_unique<Future>(req, std::move(msg));
 }
 
@@ -167,6 +169,7 @@ std::unique_ptr<Communicator::Future> MPI::recv(
     mpi_recv_impl(
         rank, tag, recv_buffer->exclusive_data_access(), recv_buffer->size, comm_, &req
     );
+    statistics_->record_recv(recv_buffer->mem_type(), recv_buffer->size);
     return std::make_unique<Future>(req, std::move(recv_buffer));
 }
 
@@ -180,6 +183,7 @@ std::unique_ptr<Communicator::Future> MPI::recv_sync_host_data(
     );
     MPI_Request req;
     mpi_recv_impl(rank, tag, synced_buffer->data(), synced_buffer->size(), comm_, &req);
+    statistics_->record_recv(MemoryType::HOST, synced_buffer->size());
     return std::make_unique<Future>(req, std::move(synced_buffer));
 }
 
@@ -211,6 +215,7 @@ std::pair<std::unique_ptr<std::vector<std::uint8_t>>, Rank> MPI::recv_any(Tag ta
         safe_cast<std::size_t>(size) == msg->size(),
         "incorrect size of the MPI_Recv message"
     );
+    statistics_->record_recv(MemoryType::HOST, msg->size());
     return {std::move(msg), probe_status.MPI_SOURCE};
 }
 
@@ -240,6 +245,7 @@ std::unique_ptr<std::vector<std::uint8_t>> MPI::recv_from(Rank src, Tag tag) {
         safe_cast<std::size_t>(size) == msg->size(),
         "incorrect size of the MPI_Recv message"
     );
+    statistics_->record_recv(MemoryType::HOST, msg->size());
     return msg;
 }
 

cpp/src/communicator/ucxx.cpp

@@ -1221,6 +1221,7 @@ std::unique_ptr<Communicator::Future> UCXX::send(
         msg->size(),
         tag_with_rank(shared_resources_->rank(), static_cast<int>(tag))
     );
+    statistics_->record_send(MemoryType::HOST, msg->size());
     return std::make_unique<Future>(req, std::move(msg));
 }
 
@@ -1232,6 +1233,7 @@ std::unique_ptr<Communicator::Future> UCXX::send(
     auto req = get_endpoint(rank)->tagSend(
         msg->data(), msg->size, tag_with_rank(shared_resources_->rank(), tag)
     );
+    statistics_->record_send(msg->mem_type(), msg->size);
     return std::make_unique<Future>(req, std::move(msg));
 }
 
@@ -1250,6 +1252,7 @@ std::unique_ptr<Communicator::Future> UCXX::recv(
         tag_with_rank(rank, tag),
         ::ucxx::TagMaskFull
     );
+    statistics_->record_recv(recv_buffer->mem_type(), recv_buffer->size);
     return std::make_unique<Future>(req, std::move(recv_buffer));
 }
 
@@ -1265,6 +1268,7 @@ std::unique_ptr<Communicator::Future> UCXX::recv_sync_host_data(
         tag_with_rank(rank, tag),
         ::ucxx::TagMaskFull
     );
+    statistics_->record_recv(MemoryType::HOST, synced_buffer->size());
     return std::make_unique<Future>(req, std::move(synced_buffer));
 }
 
@@ -1290,6 +1294,7 @@ std::pair<std::unique_ptr<std::vector<std::uint8_t>>, Rank> UCXX::recv_any(Tag t
     }
     req->checkError();
 
+    statistics_->record_recv(MemoryType::HOST, msg->size());
     return {std::move(msg), sender_rank};
 }
 
@@ -1314,6 +1319,7 @@ std::unique_ptr<std::vector<std::uint8_t>> UCXX::recv_from(Rank src, Tag tag) {
     }
     req->checkError();
 
+    statistics_->record_recv(MemoryType::HOST, msg->size());
     return msg;
 }
 

cpp/src/statistics.cpp

@@ -742,4 +742,26 @@ void Statistics::record_alloc(
     );
 }
 
+void Statistics::record_send(MemoryType src, std::size_t nbytes) {
+    static auto const names = [] {
+        std::array<std::string, MEMORY_TYPE_NAMES.size()> ret;
+        std::ranges::transform(MEMORY_TYPE_NAMES, ret.begin(), [](char const* n) {
+            return std::format("send-from-{}", n);
+        });
+        return ret;
+    }();
+    add_stat(names[static_cast<std::size_t>(src)], static_cast<double>(nbytes));
+}
+
+void Statistics::record_recv(MemoryType dst, std::size_t nbytes) {
+    static auto const names = [] {
+        std::array<std::string, MEMORY_TYPE_NAMES.size()> ret;
+        std::ranges::transform(MEMORY_TYPE_NAMES, ret.begin(), [](char const* n) {
+            return std::format("recv-to-{}", n);
+        });
+        return ret;
+    }();
+    add_stat(names[static_cast<std::size_t>(dst)], static_cast<double>(nbytes));
+}
+
 }  // namespace rapidsmpf

cpp/tests/main/mpi.cpp

@@ -27,11 +27,12 @@ void Environment::SetUp() {
 
     options_ = rapidsmpf::config::Options(rapidsmpf::config::get_environment_variables());
 
+    // enable statistics globally in the communicator
     comm_ = std::make_shared<rapidsmpf::MPI>(
         mpi_comm_,
         options_,
         std::make_shared<rapidsmpf::ProgressThread>(
-            rapidsmpf::Statistics::from_options(options_)
+            rapidsmpf::Statistics::create(rapidsmpf::Statistics::Mode::Enabled)
         )
     );
 }

cpp/tests/main/ucxx.cpp

@@ -36,11 +36,12 @@ void Environment::SetUp() {
     );
 
     options_ = rapidsmpf::config::Options(rapidsmpf::config::get_environment_variables());
+    // enable statistics globally in the communicator
     comm_ = rapidsmpf::ucxx::init_using_mpi(
         MPI_COMM_WORLD,
         options_,
         std::make_shared<rapidsmpf::ProgressThread>(
-            rapidsmpf::Statistics::from_options(options_)
+            rapidsmpf::Statistics::create(rapidsmpf::Statistics::Mode::Enabled)
         )
     );
 }

cpp/tests/test_allgather.cpp

@@ -271,6 +271,67 @@ TEST_P(AllGatherOrderedTest, non_uniform_inserts) {
     }
 }
 
+// Test that send/recv statistics record correct byte counts when allgather is run with an
+// enabled Statistics object.
+//
+// In the ring allgather each data chunk travels (nranks-1) hops. Every forwarding rank
+// records one DEVICE send and the receiving rank records one DEVICE recv, so the
+// aggregate per-rank totals are both (nranks-1) * data_bytes.
+TEST_F(BaseAllGatherTest, stats_egress_ingress) {
+    auto const& comm = GlobalEnvironment->comm_;
+    if (comm->nranks() == 1) {
+        GTEST_SKIP()
+            << "Stats test requires multiple ranks (no network traffic on 1 rank)";
+    }
+
+    constexpr int n_elements = 10;
+    constexpr std::size_t data_bytes = n_elements * sizeof(int);
+
+    // Note: there is split-brain scenario here because communicator and br have their own
+    // statistics objects
+    auto stats = comm->statistics();
+    EXPECT_TRUE(stats->enabled());
+
+    stats->clear();  // clear any previous stats since communicator stats are global
+
+    AllGather allgather{comm, 0, br.get()};
+
+    allgather.insert(
+        0, generate_packed_data(n_elements, comm->rank() * n_elements, stream, *br)
+    );
+    allgather.insert_finished();
+
+    std::vector<rapidsmpf::PackedData> results;
+    EXPECT_NO_THROW(
+        results =
+            allgather.wait_and_extract(AllGather::Ordered::NO, std::chrono::seconds{30})
+    );
+
+    auto const nranks = comm->nranks();
+    ASSERT_EQ(results.size(), static_cast<std::size_t>(nranks));
+    for (auto const& result : results) {
+        EXPECT_EQ(result.data->size, data_bytes);
+    }
+    auto const metadata_bytes = results[0].metadata->size();
+
+    // DEVICE: exact — each chunk does (nranks-1) hops
+    auto const expected_device_bytes = static_cast<double>((nranks - 1) * data_bytes);
+    EXPECT_DOUBLE_EQ(stats->get_stat("send-from-DEVICE").value(), expected_device_bytes);
+    EXPECT_DOUBLE_EQ(stats->get_stat("recv-to-DEVICE").value(), expected_device_bytes);
+
+    // HOST: lower bound — at minimum (nranks-1) data-chunk sends/recvs, each carrying
+    // at least metadata_bytes of user content (ChunkID and data_size fields add more).
+    auto const host_lower_bound = static_cast<double>((nranks - 1) * metadata_bytes);
+    auto const sent_host_bytes = stats->get_stat("send-from-HOST").value();
+    auto const recv_host_bytes = stats->get_stat("recv-to-HOST").value();
+    // metadata sends and recieves are symmetric
+    EXPECT_GT(sent_host_bytes, host_lower_bound);
+    EXPECT_GT(recv_host_bytes, host_lower_bound);
+    EXPECT_DOUBLE_EQ(sent_host_bytes, recv_host_bytes);
+
+    stats->clear();
+}
+
 // Test that reusing an OpID after a completed allgather doesn't cause cross-matching of
 // messages between the old and new collective.
 //

cpp/tests/test_allreduce.cu

@@ -305,6 +305,45 @@ class AllReduceIntSumTest
         mem_type = std::get<1>(GetParam());
     }
 
+    void run_basic_allreduce_sum_int_test() {
+        auto this_rank = comm->rank();
+        auto nranks = comm->nranks();
+
+        // Choose operator based on reduction type
+        ReduceOperator kernel =
+            mem_type == MemoryType::DEVICE
+                ? rapidsmpf::coll::detail::make_device_reduce_operator<int>(SumOp<int>{})
+                : rapidsmpf::coll::detail::make_host_reduce_operator<int>(SumOp<int>{});
+
+        std::vector<int> data(std::max(0, n_elements));
+        for (int j = 0; j < n_elements; j++) {
+            data[j] = this_rank;
+        }
+
+        auto in_buffer = make_buffer<int>(br.get(), data.data(), data.size(), mem_type);
+        auto reservation = br->reserve_or_fail(in_buffer->size, in_buffer->mem_type());
+        auto out_buffer =
+            br->make_buffer(in_buffer->size, in_buffer->stream(), reservation);
+
+        AllReduce allreduce(
+            GlobalEnvironment->comm_,
+            std::move(in_buffer),
+            std::move(out_buffer),
+            OpID{0},
+            std::move(kernel)
+        );
+
+        auto [in_result, out_result] = allreduce.wait_and_extract();
+
+        auto reduced = unpack_to_host<int>(*out_result);
+        ASSERT_EQ(static_cast<std::size_t>(n_elements), reduced.size());
+        // Expected value is sum of all ranks (0 + 1 + 2 + ... + nranks-1)
+        int const expected_value = (nranks * (nranks - 1)) / 2;
+        EXPECT_THAT(reduced, ::testing::Each(expected_value));
+
+        EXPECT_TRUE(allreduce.finished());
+    }
+
     int n_elements{};
     MemoryType mem_type{};
 };
@@ -329,41 +368,67 @@ INSTANTIATE_TEST_SUITE_P(
 );
 
 TEST_P(AllReduceIntSumTest, basic_allreduce_sum_int) {
-    auto this_rank = comm->rank();
-    auto nranks = comm->nranks();
-
-    // Choose operator based on reduction type
-    ReduceOperator kernel =
-        mem_type == MemoryType::DEVICE
-            ? rapidsmpf::coll::detail::make_device_reduce_operator<int>(SumOp<int>{})
-            : rapidsmpf::coll::detail::make_host_reduce_operator<int>(SumOp<int>{});
+    EXPECT_NO_THROW(run_basic_allreduce_sum_int_test());
+}
 
-    std::vector<int> data(std::max(0, n_elements));
-    for (int j = 0; j < n_elements; j++) {
-        data[j] = this_rank;
+TEST_P(AllReduceIntSumTest, stats_egress_ingress) {
+    if (comm->nranks() == 1) {
+        GTEST_SKIP()
+            << "Stats test requires multiple ranks (no network traffic on 1 rank)";
     }
 
-    auto in_buffer = make_buffer<int>(br.get(), data.data(), data.size(), mem_type);
-    auto reservation = br->reserve_or_fail(in_buffer->size, in_buffer->mem_type());
-    auto out_buffer = br->make_buffer(in_buffer->size, in_buffer->stream(), reservation);
+    auto stats = comm->statistics();
+    stats->clear();  // clear any previous stats since communicator stats are global
 
-    AllReduce allreduce(
-        GlobalEnvironment->comm_,
-        std::move(in_buffer),
-        std::move(out_buffer),
-        OpID{0},
-        std::move(kernel)
-    );
+    EXPECT_NO_THROW(run_basic_allreduce_sum_int_test());
 
-    auto [in_result, out_result] = allreduce.wait_and_extract();
+    auto const rank = comm->rank();
+    auto const nranks = comm->nranks();
+    auto const data_bytes = static_cast<double>(n_elements * sizeof(int));
+
+    // Compute the butterfly round count and pre-remainder size.
+    // nearest_pow2 = largest power of 2 <= nranks
+    // remainder    = nranks - nearest_pow2
+    // butterfly_rounds = log2(nearest_pow2)
+    auto const nearest_pow2 =
+        static_cast<int>(std::bit_floor(static_cast<unsigned>(nranks)));
+    auto const remainder = nranks - nearest_pow2;
+    // bit_width of a power-of-2 n is log2(n)+1, so subtract 1 to get log2.
+    auto const butterfly_rounds =
+        static_cast<int>(std::bit_width(static_cast<unsigned>(nearest_pow2))) - 1;
+
+    // Expected traffic per rank:
+    //   Even pre-remainder (rank < 2*r, even): 1 pre-send + 1 post-recv  → 1×D each
+    //   Odd  pre-remainder (rank < 2*r, odd):  1 pre-recv + B butterfly
+    //                                           + 1 post-send             → (1+B)×D each
+    //   Non-remainder      (rank >= 2*r):       B butterfly rounds        → B×D each
+    double expected_egress{};
+    double expected_ingress{};
+    if (rank < 2 * remainder) {
+        if (rank % 2 == 0) {
+            expected_egress = data_bytes;
+            expected_ingress = data_bytes;
+        } else {
+            expected_egress = (1 + butterfly_rounds) * data_bytes;
+            expected_ingress = (1 + butterfly_rounds) * data_bytes;
+        }
+    } else {
+        expected_egress = butterfly_rounds * data_bytes;
+        expected_ingress = butterfly_rounds * data_bytes;
+    }
 
-    auto reduced = unpack_to_host<int>(*out_result);
-    ASSERT_EQ(static_cast<std::size_t>(n_elements), reduced.size());
-    // Expected value is sum of all ranks (0 + 1 + 2 + ... + nranks-1)
-    int const expected_value = (nranks * (nranks - 1)) / 2;
-    EXPECT_THAT(reduced, ::testing::Each(expected_value));
+    EXPECT_DOUBLE_EQ(
+        stats->get_stat(std::format("send-from-{}", rapidsmpf::to_string(mem_type)))
+            .value(),
+        expected_egress
+    );
+    EXPECT_DOUBLE_EQ(
+        stats->get_stat(std::format("recv-to-{}", rapidsmpf::to_string(mem_type)))
+            .value(),
+        expected_ingress
+    );
 
-    EXPECT_TRUE(allreduce.finished());
+    stats->clear();
 }
 
 template <typename T, typename Op, MemoryType MemType>