Move AllGather/AllReduce to coll namespace

Author: pentschev

cpp/CMakeLists.txt

@@ -156,9 +156,9 @@ target_link_options(maybe_asan INTERFACE "$<$<BOOL:${RAPIDSMPF_ASAN}>:-fsanitize
 
 add_library(
   rapidsmpf
-  src/allgather/allgather.cpp
-  src/allreduce/allreduce.cpp
-  src/allreduce/device_kernels.cu
+  src/coll/allgather.cpp
+  src/coll/allreduce.cpp
+  src/coll/device_kernels.cu
   src/bootstrap/bootstrap.cpp
   src/bootstrap/file_backend.cpp
   src/communicator/communicator.cpp

cpp/include/rapidsmpf/allgather/allgather.hpp renamed to cpp/include/rapidsmpf/coll/allgather.hpp

@@ -27,13 +27,13 @@
 #include <rapidsmpf/statistics.hpp>
 
 /**
- * @namespace rapidsmpf::allgather
- * @brief Allgather communication interfaces.
+ * @namespace rapidsmpf::coll
+ * @brief Collective communication interfaces.
  *
  * An allgather service for distributed communication where all ranks collect
  * data from all other ranks.
  */
-namespace rapidsmpf::allgather {
+namespace rapidsmpf::coll {
 namespace detail {
 
 /// @brief Type alias for chunk identifiers.
@@ -559,4 +559,4 @@ class AllGather {
     std::vector<std::unique_ptr<Communicator::Future>> receive_futures_{};
 };
 
-}  // namespace rapidsmpf::allgather
+}  // namespace rapidsmpf::coll

cpp/include/rapidsmpf/allreduce/allreduce.hpp renamed to cpp/include/rapidsmpf/coll/allreduce.hpp

@@ -15,7 +15,7 @@
 #include <utility>
 #include <vector>
 
-#include <rapidsmpf/allgather/allgather.hpp>
+#include <rapidsmpf/coll/allgather.hpp>
 #include <rapidsmpf/communicator/communicator.hpp>
 #include <rapidsmpf/error.hpp>
 #include <rapidsmpf/memory/buffer.hpp>
@@ -24,7 +24,7 @@
 #include <rapidsmpf/progress_thread.hpp>
 #include <rapidsmpf/statistics.hpp>
 
-namespace rapidsmpf::allreduce {
+namespace rapidsmpf::coll {
 
 /**
  * @brief Reduction operators supported by `AllReduce`.
@@ -57,7 +57,7 @@ using ReduceKernel = std::function<void(PackedData& accum, PackedData&& incoming
 /**
  * @brief AllReduce collective.
  *
- * The current implementation is built using `allgather::AllGather` and performs
+ * The current implementation is built using `coll::AllGather` and performs
  * the reduction locally after allgather completes. Considering `R` is the number of
  * ranks, and `N` is the number of bytes of data, per rank this incurs `O(R * N)` bytes of
  * memory consumption and `O(R)` communication operations.
@@ -89,7 +89,7 @@ class AllReduce {
      * @param finished_callback Optional callback run once locally when the allreduce
      *        is finished and results are ready for extraction.
      *
-     * @note This constructor internally creates an `allgather::AllGather` instance
+     * @note This constructor internally creates an `AllGather` instance
      *       that uses the same communicator, progress thread, and buffer resource.
      */
     AllReduce(
@@ -183,7 +183,7 @@ class AllReduce {
     ReduceKernel reduce_kernel_;  ///< Type-erased reduction kernel
     std::function<void(void)> finished_callback_;  ///< Optional finished callback
 
-    allgather::AllGather gatherer_;  ///< Underlying allgather primitive
+    AllGather gatherer_;  ///< Underlying allgather primitive
 
     std::atomic<std::uint32_t> nlocal_insertions_{0};  ///< Number of local inserts
     std::atomic<bool> reduced_computed_{
@@ -204,4 +204,4 @@ template <typename T, ReduceOp Op>
 ReduceKernel make_reduce_kernel();
 }  // namespace detail
 
-}  // namespace rapidsmpf::allreduce
+}  // namespace rapidsmpf::coll

cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp

@@ -31,8 +31,10 @@ namespace rapidsmpf {
 /**
  * @brief Checks if the PinnedMemoryResource is supported for the current CUDA version.
  *
- * Requires rapidsmpf to be build with cuda>=12.6.
+ * Requires rapidsmpf to be built with CUDA>=12.6.
  *
+ * @return True if the PinnedMemoryResource is supported for the current CUDA version,
+ *         false otherwise.
  * @note The driver version check is cached and only performed once.
  */
 inline bool is_pinned_memory_resources_supported() {

cpp/include/rapidsmpf/streaming/coll/allgather.hpp

@@ -9,7 +9,7 @@
 #include <memory>
 #include <vector>
 
-#include <rapidsmpf/allgather/allgather.hpp>
+#include <rapidsmpf/coll/allgather.hpp>
 #include <rapidsmpf/communicator/communicator.hpp>
 #include <rapidsmpf/memory/packed_data.hpp>
 #include <rapidsmpf/streaming/chunks/packed_data.hpp>
@@ -22,7 +22,7 @@
 namespace rapidsmpf::streaming {
 
 /**
- * @brief Asynchronous (coroutine) interface to `allgather::AllGather`.
+ * @brief Asynchronous (coroutine) interface to `coll::AllGather`.
  *
  * Once the AllGather is created, many tasks may insert data into it. If multiple tasks
  * insert data, the user is responsible for arranging that `insert_finished` is only
@@ -31,8 +31,8 @@ namespace rapidsmpf::streaming {
  */
 class AllGather {
   public:
-    /// @copydoc allgather::AllGather::Ordered
-    using Ordered = rapidsmpf::allgather::AllGather::Ordered;
+    /// @copydoc coll::AllGather::Ordered
+    using Ordered = rapidsmpf::coll::AllGather::Ordered;
     /**
      * @brief Construct an asynchronous allgather.
      *
@@ -63,7 +63,7 @@ class AllGather {
      */
     void insert(std::uint64_t sequence_number, PackedDataChunk&& chunk);
 
-    /// @copydoc rapidsmpf::allgather::AllGather::insert_finished()
+    /// @copydoc rapidsmpf::coll::AllGather::insert_finished()
     void insert_finished();
 
     /**
@@ -82,15 +82,15 @@ class AllGather {
     coro::event
         event_{};  ///< Event tracking whether all data has arrived and can be extracted.
     std::shared_ptr<Context> ctx_;  ///< Streaming context.
-    allgather::AllGather gatherer_;  ///< Underlying collective allgather.
+    coll::AllGather gatherer_;  ///< Underlying collective allgather.
 };
 
 namespace node {
 
 /**
  * @brief Create an allgather node for a single allgather operation.
  *
- * This is a streaming version of `rapidsmpf::allgather::AllGather` that operates on
+ * This is a streaming version of `rapidsmpf::coll::AllGather` that operates on
  * packed data received through `Channel`s.
  *
  * @param ctx The streaming context to use.

cpp/src/allgather/allgather.cpp renamed to cpp/src/coll/allgather.cpp

@@ -12,13 +12,13 @@
 #include <mutex>
 #include <optional>
 
-#include <rapidsmpf/allgather/allgather.hpp>
+#include <rapidsmpf/coll/allgather.hpp>
 #include <rapidsmpf/communicator/communicator.hpp>
 #include <rapidsmpf/memory/buffer.hpp>
 #include <rapidsmpf/progress_thread.hpp>
 #include <rapidsmpf/utils.hpp>
 
-namespace rapidsmpf::allgather {
+namespace rapidsmpf::coll {
 namespace detail {
 
 Chunk::Chunk(
@@ -538,4 +538,4 @@ ProgressThread::ProgressState AllGather::event_loop() {
                    : ProgressThread::ProgressState::InProgress;
 }
 
-}  // namespace rapidsmpf::allgather
+}  // namespace rapidsmpf::coll

cpp/src/allreduce/allreduce.cpp renamed to cpp/src/coll/allreduce.cpp

@@ -10,10 +10,10 @@
 #include <utility>
 #include <vector>
 
-#include <rapidsmpf/allreduce/allreduce.hpp>
+#include <rapidsmpf/coll/allreduce.hpp>
 #include <rapidsmpf/error.hpp>
 
-namespace rapidsmpf::allreduce {
+namespace rapidsmpf::coll {
 
 AllReduce::AllReduce(
     std::shared_ptr<Communicator> comm,
@@ -56,9 +56,8 @@ std::vector<PackedData> AllReduce::wait_and_extract(std::chrono::milliseconds ti
     // Block until the underlying allgather completes, then perform the reduction locally
     // (exactly once).
     if (!reduced_computed_.load(std::memory_order_acquire)) {
-        auto gathered = gatherer_.wait_and_extract(
-            allgather::AllGather::Ordered::YES, std::move(timeout)
-        );
+        auto gathered =
+            gatherer_.wait_and_extract(AllGather::Ordered::YES, std::move(timeout));
         reduced_results_ = reduce_all(std::move(gathered));
         reduced_computed_.store(true, std::memory_order_release);
         if (finished_callback_) {
@@ -120,4 +119,4 @@ std::vector<PackedData> AllReduce::reduce_all(std::vector<PackedData>&& gathered
     return results;
 }
 
-}  // namespace rapidsmpf::allreduce
+}  // namespace rapidsmpf::coll

cpp/src/allreduce/device_kernels.cu renamed to cpp/src/coll/device_kernels.cu

@@ -8,11 +8,11 @@
 
 #include <rmm/exec_policy.hpp>
 
-#include <rapidsmpf/allreduce/allreduce.hpp>
+#include <rapidsmpf/coll/allreduce.hpp>
 #include <rapidsmpf/error.hpp>
 #include <rapidsmpf/memory/buffer.hpp>
 
-namespace rapidsmpf::allreduce::detail {
+namespace rapidsmpf::coll::detail {
 
 namespace {
 
@@ -177,4 +177,4 @@ ReduceKernel make_reduce_kernel<unsigned long, ReduceOp::MAX>() {
     return make_reduce_kernel_impl<unsigned long, ReduceOp::MAX>();
 }
 
-}  // namespace rapidsmpf::allreduce::detail
+}  // namespace rapidsmpf::coll::detail

cpp/src/streaming/coll/allgather.cpp

@@ -14,7 +14,7 @@ namespace rapidsmpf::streaming {
 
 AllGather::AllGather(std::shared_ptr<Context> ctx, OpID op_id)
     : ctx_{std::move(ctx)},
-      gatherer_{allgather::AllGather(
+      gatherer_{coll::AllGather(
           ctx_->comm(),
           ctx_->progress_thread(),
           op_id,

cpp/tests/streaming/test_read_parquet.cpp

@@ -32,7 +32,7 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 
-#include <rapidsmpf/allgather/allgather.hpp>
+#include <rapidsmpf/coll/allgather.hpp>
 #include <rapidsmpf/integrations/cudf/partition.hpp>
 #include <rapidsmpf/memory/packed_data.hpp>
 #include <rapidsmpf/owning_wrapper.hpp>
@@ -228,7 +228,7 @@ TEST_P(StreamingReadParquetParams, ReadParquet) {
     }
     run_streaming_pipeline(std::move(nodes));
 
-    allgather::AllGather allgather(
+    coll::AllGather allgather(
         GlobalEnvironment->comm_,
         GlobalEnvironment->progress_thread_,
         /* op_id = */ 0,
@@ -254,8 +254,7 @@ TEST_P(StreamingReadParquetParams, ReadParquet) {
     allgather.insert_finished();
 
     // May as well check on all ranks, so we also mildly exercise the allgather.
-    auto gathered_packed_data =
-        allgather.wait_and_extract(allgather::AllGather::Ordered::YES);
+    auto gathered_packed_data = allgather.wait_and_extract(coll::AllGather::Ordered::YES);
     auto result = unpack_and_concat(
         std::move(gathered_packed_data), rmm::cuda_stream_default, br.get()
     );