feat: add collective communication for mccl

CMakeLists.txt

@@ -1,5 +1,6 @@
 cmake_minimum_required(VERSION 3.18)
 project(InfiniCCL VERSION 0.1.0 LANGUAGES C CXX)
+include(CTest)
 
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
@@ -23,6 +24,9 @@ option(WITH_OMPI "Enable OpenMPI backend" OFF)
 option(WITH_MPICH "Enable MPICH backend" OFF)
 option(WITH_NCCL "Enable NCCL backend"    OFF)
 option(WITH_MCCL "Enable MCCL backend"    OFF)
+if(BUILD_TESTING)
+  add_subdirectory(tests)
+endif()
 
 # =========================================================
 # --- MISC. BUILD OPTIONS ---

examples/ccl/collectives.cc

@@ -0,0 +1,160 @@
+/**
+ * InfiniCCL Example: Thread-per-GPU Single-Node MCCL Broadcast
+ *
+ * This example demonstrates spawning one CPU thread per GPU on a single node,
+ * generating a shared UniqueID, and performing a Broadcast entirely via
+ * InfiniCCL's native CCL backend.
+ *
+ * Note: to properly run this example, you should specify `--launcher none`
+ * since it requires no MPI process spawning.
+ */
+
+#include <unistd.h>
+
+#include <iostream>
+#include <numeric>
+#include <thread>
+#include <vector>
+
+// Public API
+#include "infiniccl.h"
+
+// Example-Specific Utilities
+#include "utils.h"
+
+// Internal Headers (Accessible via example-specific include paths, technically
+// not public APIs)
+#include "backend_manifest.h"
+
+using namespace infini::ccl;
+
+// Structure to pass execution data to each GPU worker thread.
+struct ThreadArgs {
+  int rank;
+  int size;
+  infinicclUniqueId id;
+  size_t num_elements;
+  int warmup_iter;
+  int profile_iter;
+};
+
+// Worker function executed by each CPU thread.
+void WorkerThread(ThreadArgs args) {
+  constexpr Device::Type kDevType =
+      ListGetBest<DevicePriority>(EnabledDevices{});
+  using Rt = Runtime<kDevType>;
+
+  // Bind this specific CPU thread to its designated local GPU device
+  // In a thread-per-GPU model, `local_rank` is exactly the thread's rank ID.
+  int local_device_id = args.rank;
+  CHECK_RT(Rt, Rt::SetDevice(local_device_id));
+
+  infinicclComm_t comm = nullptr;
+  CHECK_INFINI(infinicclCommInitRank(&comm, args.size, args.id, args.rank));
+
+  // Prepare Host Data Structures
+  std::vector<float> h_send(args.num_elements);
+  std::vector<float> h_recv(args.num_elements, 0.0f);
+
+  // Each rank provides its own (rank + 1) as data.
+  for (size_t i = 0; i < args.num_elements; i++) {
+    h_send[i] = static_cast<float>(args.rank + 1);
+  }
+
+  // Allocate GPU Memory using InfiniCCL's Runtime abstraction layer.
+  float *d_send = nullptr;
+  float *d_recv = nullptr;
+  size_t total_bytes = args.num_elements * sizeof(float);
+
+  CHECK_RT(Rt, Rt::Malloc((void **)&d_send, total_bytes));
+  CHECK_RT(Rt, Rt::Malloc((void **)&d_recv, total_bytes));
+
+  CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes,
+                          Rt::MemcpyHostToDevice));
+  CHECK_RT(Rt, Rt::Memcpy(d_recv, h_recv.data(), total_bytes,
+                          Rt::MemcpyHostToDevice));
+
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+
+  // Warm-up Iterations
+  for (int i = 0; i < args.warmup_iter; ++i) {
+    CHECK_INFINI(infinicclBroadcast(d_send, d_recv, args.num_elements,
+                                    infinicclFloat32, 0, comm,
+                                    nullptr));
+  }
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+
+  // Profiling Iterations
+  Timer timer;
+  for (int i = 0; i < args.profile_iter; i++) {
+    CHECK_INFINI(infinicclBroadcast(d_send, d_recv, args.num_elements,
+                                    infinicclFloat32, 0, comm,
+                                    nullptr));
+  }
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+  double elapsed = timer.ElapsedMs() / static_cast<double>(args.profile_iter);
+
+  // Copy broadcast data back from device to host.
+  CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes,
+                          Rt::MemcpyDeviceToHost));
+
+  // Result Validation
+  constexpr float expected = 1.0f;
+  Validator::ValidateResult(h_recv.data(), args.num_elements, expected,
+                            args.rank, true, "Broadcast");
+
+  // Metrics Reporting (Only Rank 0)
+  if (args.rank == 0) {
+    std::cout << "\n=== Single-Node Threaded MCCL Broadcast Results ==="
+              << std::endl;
+    std::cout << "Data size: " << args.num_elements << " floats ("
+              << total_bytes / 1024 / 1024 << " MB)" << std::endl;
+    Metrics metrics{elapsed, total_bytes, args.size};
+    metrics.Print();
+  }
+
+  // Cleanup local rank resources.
+  CHECK_RT(Rt, Rt::Free(d_send));
+  CHECK_RT(Rt, Rt::Free(d_recv));
+  CHECK_INFINI(infinicclCommDestroy(comm));
+}
+
+int main(int argc, char **argv) {
+  int num_gpus = 8;
+  int warmup_iters = 1;
+  int profile_iters = 20;
+  size_t num_elements = 1 << 25;
+
+  (void)argc;
+  (void)argv;
+
+  char hostname[256];
+  gethostname(hostname, sizeof(hostname));
+  std::cout << "[Main Process] Host: " << hostname
+            << " | Target GPUs: " << num_gpus << std::endl;
+
+  infinicclUniqueId shared_id;
+  CHECK_INFINI(infinicclGetUniqueId(&shared_id));
+
+  // Spawn CPU thread pool.
+  std::vector<std::thread> threads;
+  threads.reserve(num_gpus);
+
+  for (int rank = 0; rank < num_gpus; ++rank) {
+    ThreadArgs args{rank,         num_gpus,     shared_id,
+                    num_elements, warmup_iters, profile_iters};
+    threads.emplace_back(WorkerThread, args);
+  }
+
+  // Await execution completion across all threads.
+  for (auto &t : threads) {
+    if (t.joinable()) {
+      t.join();
+    }
+  }
+
+  std::cout
+      << "[Main Process] All worker threads joined. InfiniCCL finalized safely."
+      << std::endl;
+  return EXIT_SUCCESS;
+}

src/backends/ccl/common/impl/collectives.h

@@ -0,0 +1,161 @@
+#ifndef INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_COLLECTIVES_H_
+#define INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_COLLECTIVES_H_
+
+#include "backends/ccl/common/api.h"
+#include "device.h"
+#include "traits.h"
+#include "data_type_impl.h"
+#include "return_status_impl.h"
+#include "base/all_gather.h"
+#include "base/all_reduce.h"
+#include "base/all_to_all.h"
+#include "base/broadcast.h"
+#include "base/gather.h"
+#include "base/reduce.h"
+#include "base/reduce_scatter.h"
+#include "base/scatter.h"
+#include "backends/ccl/common/comm_instance.h"
+#include "communicator.h"
+
+namespace infini::ccl {
+
+template <BackendType backend, Device::Type device>
+struct CclCollective {
+  using Api = CclApi<backend, device>;
+  using TypeMap = CclTypeMap<backend, device>;
+  using CommInstance = CclCommInstance<Api>;
+
+  static CommInstance *Get(Communicator *comm) {
+    auto *instance = static_cast<CommInstance *>(comm->intra_comm());
+    return instance && instance->handle ? instance : nullptr;
+  }
+
+  static bool ToDataType(DataType value, typename Api::DataType *result) {
+    return TypeMap::ToBackendDataType(value, result);
+  }
+
+  static bool ToReduction(ReductionOpType value, typename Api::RedOp *result) {
+    return TypeMap::ToBackendRedOp(value, result);
+  }
+};
+
+template <BackendType backend, Device::Type device>
+struct CclBroadcastImpl {
+  static ReturnStatus Apply(const void *send_buff, void *recv_buff, size_t count,
+                            DataType type, int root, Communicator *comm,
+                            void *stream) {
+    using C = CclCollective<backend, device>;
+    auto *instance = C::Get(comm);
+    typename C::Api::DataType backend_type{};
+    if (!instance) return ReturnStatus::kInternalError;
+    if (!C::ToDataType(type, &backend_type)) return ReturnStatus::kNotSupported;
+    return C::Api::Check(C::Api::Broadcast(send_buff, recv_buff, count,
+                                            backend_type, root, instance->handle,
+                                            reinterpret_cast<typename C::Api::Stream>(stream)));
+  }
+};
+
+template <BackendType backend, Device::Type device>
+struct CclReduceImpl {
+  static ReturnStatus Apply(const void *send_buff, void *recv_buff, size_t count,
+                            DataType type, ReductionOpType op, int root,
+                            Communicator *comm, void *stream) {
+    using C = CclCollective<backend, device>;
+    auto *instance = C::Get(comm);
+    typename C::Api::DataType backend_type{};
+    typename C::Api::RedOp backend_op{};
+    if (!instance) return ReturnStatus::kInternalError;
+    if (!C::ToDataType(type, &backend_type) || !C::ToReduction(op, &backend_op)) {
+      return ReturnStatus::kNotSupported;
+    }
+    return C::Api::Check(C::Api::Reduce(send_buff, recv_buff, count, backend_type,
+                                        backend_op, root, instance->handle,
+                                        reinterpret_cast<typename C::Api::Stream>(stream)));
+  }
+};
+
+template <BackendType backend, Device::Type device>
+struct CclReduceScatterImpl {
+  static ReturnStatus Apply(const void *send_buff, void *recv_buff, size_t count,
+                            DataType type, ReductionOpType op, Communicator *comm,
+                            void *stream) {
+    using C = CclCollective<backend, device>;
+    auto *instance = C::Get(comm);
+    typename C::Api::DataType backend_type{};
+    typename C::Api::RedOp backend_op{};
+    if (!instance) return ReturnStatus::kInternalError;
+    if (!C::ToDataType(type, &backend_type) || !C::ToReduction(op, &backend_op)) {
+      return ReturnStatus::kNotSupported;
+    }
+    return C::Api::Check(C::Api::ReduceScatter(send_buff, recv_buff, count,
+                                               backend_type, backend_op,
+                                               instance->handle,
+                                               reinterpret_cast<typename C::Api::Stream>(stream)));
+  }
+};
+
+template <BackendType backend, Device::Type device>
+struct CclAllGatherImpl {
+  static ReturnStatus Apply(const void *send_buff, void *recv_buff, size_t count,
+                            DataType type, Communicator *comm, void *stream) {
+    using C = CclCollective<backend, device>;
+    auto *instance = C::Get(comm);
+    typename C::Api::DataType backend_type{};
+    if (!instance) return ReturnStatus::kInternalError;
+    if (!C::ToDataType(type, &backend_type)) return ReturnStatus::kNotSupported;
+    return C::Api::Check(C::Api::AllGather(send_buff, recv_buff, count,
+                                            backend_type, instance->handle,
+                                            reinterpret_cast<typename C::Api::Stream>(stream)));
+  }
+};
+
+template <BackendType backend, Device::Type device>
+struct CclGatherImpl {
+  static ReturnStatus Apply(const void *send_buff, void *recv_buff, size_t count,
+                            DataType type, int root, Communicator *comm,
+                            void *stream) {
+    using C = CclCollective<backend, device>;
+    auto *instance = C::Get(comm);
+    typename C::Api::DataType backend_type{};
+    if (!instance) return ReturnStatus::kInternalError;
+    if (!C::ToDataType(type, &backend_type)) return ReturnStatus::kNotSupported;
+    return C::Api::Check(C::Api::Gather(send_buff, recv_buff, count, backend_type,
+                                         root, instance->handle,
+                                         reinterpret_cast<typename C::Api::Stream>(stream)));
+  }
+};
+
+template <BackendType backend, Device::Type device>
+struct CclScatterImpl {
+  static ReturnStatus Apply(const void *send_buff, void *recv_buff, size_t count,
+                            DataType type, int root, Communicator *comm,
+                            void *stream) {
+    using C = CclCollective<backend, device>;
+    auto *instance = C::Get(comm);
+    typename C::Api::DataType backend_type{};
+    if (!instance) return ReturnStatus::kInternalError;
+    if (!C::ToDataType(type, &backend_type)) return ReturnStatus::kNotSupported;
+    return C::Api::Check(C::Api::Scatter(send_buff, recv_buff, count, backend_type,
+                                          root, instance->handle,
+                                          reinterpret_cast<typename C::Api::Stream>(stream)));
+  }
+};
+
+template <BackendType backend, Device::Type device>
+struct CclAllToAllImpl {
+  static ReturnStatus Apply(const void *send_buff, void *recv_buff, size_t count,
+                            DataType type, Communicator *comm, void *stream) {
+    using C = CclCollective<backend, device>;
+    auto *instance = C::Get(comm);
+    typename C::Api::DataType backend_type{};
+    if (!instance) return ReturnStatus::kInternalError;
+    if (!C::ToDataType(type, &backend_type)) return ReturnStatus::kNotSupported;
+    return C::Api::Check(C::Api::AllToAll(send_buff, recv_buff, count,
+                                           backend_type, instance->handle,
+                                           reinterpret_cast<typename C::Api::Stream>(stream)));
+  }
+};
+
+}  // namespace infini::ccl
+
+#endif  // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_COLLECTIVES_H_

src/backends/ccl/mccl/api.h

@@ -49,6 +49,51 @@ struct McclApi {
     return mcclAllReduce(send_buff, recv_buff, count, data_type, op, comm,
                          stream);
   }
+  static Result Broadcast(const void *send_buff, void *recv_buff, size_t count,
+                          DataType data_type, int root, Comm comm,
+                          Stream stream) {
+    return mcclBroadcast(send_buff, recv_buff, count, data_type, root, comm,
+                         stream);
+  }
+
+  static Result Reduce(const void *send_buff, void *recv_buff, size_t count,
+                       DataType data_type, RedOp op, int root, Comm comm,
+                       Stream stream) {
+    return mcclReduce(send_buff, recv_buff, count, data_type, op, root, comm,
+                      stream);
+  }
+
+  static Result ReduceScatter(const void *send_buff, void *recv_buff,
+                              size_t recv_count, DataType data_type, RedOp op,
+                              Comm comm, Stream stream) {
+    return mcclReduceScatter(send_buff, recv_buff, recv_count, data_type, op,
+                             comm, stream);
+  }
+
+  static Result AllGather(const void *send_buff, void *recv_buff,
+                          size_t send_count, DataType data_type, Comm comm,
+                          Stream stream) {
+    return mcclAllGather(send_buff, recv_buff, send_count, data_type, comm,
+                         stream);
+  }
+
+  static Result Gather(const void *send_buff, void *recv_buff, size_t send_count,
+                       DataType data_type, int root, Comm comm, Stream stream) {
+    return mcclGather(send_buff, recv_buff, send_count, data_type, root, comm,
+                      stream);
+  }
+
+  static Result Scatter(const void *send_buff, void *recv_buff,
+                        size_t recv_count, DataType data_type, int root,
+                        Comm comm, Stream stream) {
+    return mcclScatter(send_buff, recv_buff, recv_count, data_type, root, comm,
+                       stream);
+  }
+
+  static Result AllToAll(const void *send_buff, void *recv_buff, size_t count,
+                         DataType data_type, Comm comm, Stream stream) {
+    return mcclAllToAll(send_buff, recv_buff, count, data_type, comm, stream);
+  }
 };
 
 }  // namespace infini::ccl