feat(nccl): add all-gather communication

examples/ccl/all_gather.cc

@@ -0,0 +1,123 @@
+/**
+ * InfiniCCL Example: Thread-per-GPU Single-Node AllGather
+ *
+ * This example validates out-of-place and in-place AllGather across two GPUs
+ * through InfiniCCL's native CCL backend without an MPI launcher.
+ */
+
+#include <algorithm>
+#include <atomic>
+#include <cstddef>
+#include <cstdlib>
+#include <iostream>
+#include <thread>
+#include <vector>
+
+#include "backend_manifest.h"
+#include "infiniccl.h"
+#include "utils.h"
+
+using namespace infini::ccl;
+
+namespace {
+
+constexpr int kRankCount = 2;
+constexpr size_t kNumElements = 1 << 10;
+
+struct ThreadArgs {
+  int rank;
+  infinicclUniqueId id;
+  std::atomic_bool* all_correct;
+};
+
+bool Validate(const std::vector<float>& output, int rank, const char* mode) {
+  for (int source = 0; source < kRankCount; ++source) {
+    const float expected = static_cast<float>(source + 1);
+    const size_t offset = static_cast<size_t>(source) * kNumElements;
+    for (size_t i = 0; i < kNumElements; ++i) {
+      if (output[offset + i] != expected) {
+        std::cerr << mode << " validation failed on rank " << rank
+                  << " at source rank " << source << ", element " << i
+                  << ": expected " << expected << ", got " << output[offset + i]
+                  << "." << std::endl;
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+void WorkerThread(ThreadArgs args) {
+  constexpr Device::Type kDevType =
+      ListGetBest<DevicePriority>(EnabledDevices{});
+  using Rt = Runtime<kDevType>;
+
+  CHECK_RT(Rt, Rt::SetDevice(args.rank));
+
+  infinicclComm_t comm = nullptr;
+  CHECK_INFINI(infinicclCommInitRank(&comm, kRankCount, args.id, args.rank));
+
+  std::vector<float> host_send(kNumElements, static_cast<float>(args.rank + 1));
+  std::vector<float> host_recv(kNumElements * kRankCount, 0.0f);
+  float* device_send = nullptr;
+  float* device_recv = nullptr;
+  const size_t send_bytes = kNumElements * sizeof(float);
+  const size_t recv_bytes = send_bytes * kRankCount;
+
+  CHECK_RT(Rt, Rt::Malloc(reinterpret_cast<void**>(&device_send), send_bytes));
+  CHECK_RT(Rt, Rt::Malloc(reinterpret_cast<void**>(&device_recv), recv_bytes));
+  CHECK_RT(Rt, Rt::Memcpy(device_send, host_send.data(), send_bytes,
+                          Rt::MemcpyHostToDevice));
+
+  CHECK_INFINI(infinicclAllGather(device_send, device_recv, kNumElements,
+                                  infinicclFloat32, comm, nullptr));
+  CHECK_RT(Rt, Rt::Memcpy(host_recv.data(), device_recv, recv_bytes,
+                          Rt::MemcpyDeviceToHost));
+  if (!Validate(host_recv, args.rank, "Out-of-place AllGather")) {
+    args.all_correct->store(false, std::memory_order_relaxed);
+  }
+
+  std::fill(host_recv.begin(), host_recv.end(), 0.0f);
+  std::fill_n(host_recv.begin() + static_cast<size_t>(args.rank) * kNumElements,
+              kNumElements, static_cast<float>(args.rank + 1));
+  CHECK_RT(Rt, Rt::Memcpy(device_recv, host_recv.data(), recv_bytes,
+                          Rt::MemcpyHostToDevice));
+
+  float* local_block =
+      device_recv + static_cast<size_t>(args.rank) * kNumElements;
+  CHECK_INFINI(infinicclAllGather(local_block, device_recv, kNumElements,
+                                  infinicclFloat32, comm, nullptr));
+  CHECK_RT(Rt, Rt::Memcpy(host_recv.data(), device_recv, recv_bytes,
+                          Rt::MemcpyDeviceToHost));
+  if (!Validate(host_recv, args.rank, "In-place AllGather")) {
+    args.all_correct->store(false, std::memory_order_relaxed);
+  }
+
+  CHECK_RT(Rt, Rt::Free(device_send));
+  CHECK_RT(Rt, Rt::Free(device_recv));
+  CHECK_INFINI(infinicclCommDestroy(comm));
+}
+
+}  // namespace
+
+int main() {
+  infinicclUniqueId shared_id;
+  CHECK_INFINI(infinicclGetUniqueId(&shared_id));
+
+  std::atomic_bool all_correct{true};
+  std::vector<std::thread> threads;
+  threads.reserve(kRankCount);
+  for (int rank = 0; rank < kRankCount; ++rank) {
+    threads.emplace_back(WorkerThread,
+                         ThreadArgs{rank, shared_id, &all_correct});
+  }
+  for (auto& thread : threads) {
+    thread.join();
+  }
+
+  if (!all_correct.load(std::memory_order_relaxed)) {
+    return EXIT_FAILURE;
+  }
+  std::cout << "AllGather validation passed." << std::endl;
+  return EXIT_SUCCESS;
+}

examples/mpi/all_gather.cc

@@ -24,7 +24,7 @@
 
 using namespace infini::ccl;
 
-void RunAllGatherExample(int argc, char **argv, int warmup_iter,
+bool RunAllGatherExample(int argc, char **argv, int warmup_iter,
                          int profile_iter, const size_t kNumElements) {
   constexpr Device::Type kDevType =
       ListGetBest<DevicePriority>(EnabledDevices{});
@@ -113,14 +113,14 @@ void RunAllGatherExample(int argc, char **argv, int warmup_iter,
 
   // Result Validation
   bool correct = true;
-  int error_count = 0;
 
   for (int src_rank = 0; src_rank < size; ++src_rank) {
     float expected = static_cast<float>(src_rank + 1);
     size_t offset = static_cast<size_t>(src_rank) * kNumElements;
 
-    Validator::ValidateResult(h_recv.data() + offset, kNumElements, expected,
-                              rank);
+    correct = Validator::ValidateResult(h_recv.data() + offset, kNumElements,
+                                        expected, rank) &&
+              correct;
   }
 
   if (rank == 0) {
@@ -132,7 +132,6 @@ void RunAllGatherExample(int argc, char **argv, int warmup_iter,
     std::cout << "Correct: "
               << (correct ? (GREEN + std::string("YES") + RESET)
                           : (RED + std::string("NO") + RESET));
-    if (!correct) std::cout << " (" << error_count << " errors)";
     std::cout << std::endl;
 
     std::cout << "Sample blocks: ";
@@ -159,14 +158,16 @@ void RunAllGatherExample(int argc, char **argv, int warmup_iter,
   if (rank == 0) {
     std::cout << "InfiniCCL finalized." << std::endl;
   }
+  return correct;
 }
 
 int main(int argc, char **argv) {
   int warmup_iters = 2;
   int profile_iters = 20;
   size_t num_elements = 1 << 20;
 
-  RunAllGatherExample(argc, argv, warmup_iters, profile_iters, num_elements);
-
-  return EXIT_SUCCESS;
+  return RunAllGatherExample(argc, argv, warmup_iters, profile_iters,
+                             num_elements)
+             ? EXIT_SUCCESS
+             : EXIT_FAILURE;
 }

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

@@ -0,0 +1,44 @@
+#ifndef INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_GATHER_H_
+#define INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_GATHER_H_
+
+#include "backends/ccl/common/api.h"
+#include "backends/ccl/common/comm_instance.h"
+#include "base/all_gather.h"
+#include "communicator.h"
+
+namespace infini::ccl {
+
+template <BackendType backend, Device::Type device>
+class CclAllGatherImpl {
+ public:
+  static ReturnStatus Apply(const void* send_buff, void* recv_buff,
+                            size_t count, DataType data_type,
+                            Communicator* comm, void* stream) {
+    using Api = CclApi<backend, device>;
+    using TypeMap = CclTypeMap<backend, device>;
+    using CommInstance = CclCommInstance<Api>;
+
+    if (!comm || !comm->intra_comm() || comm->intra_comm_backend() != backend ||
+        comm->device_type() != device) {
+      return ReturnStatus::kInternalError;
+    }
+
+    auto* intra = static_cast<CommInstance*>(comm->intra_comm());
+    if (!intra->handle) {
+      return ReturnStatus::kInternalError;
+    }
+
+    typename Api::DataType ccl_type{};
+    if (!TypeMap::ToBackendDataType(data_type, &ccl_type)) {
+      return ReturnStatus::kNotSupported;
+    }
+
+    return Api::Check(
+        Api::AllGather(send_buff, recv_buff, count, ccl_type, intra->handle,
+                       reinterpret_cast<typename Api::Stream>(stream)));
+  }
+};
+
+}  // namespace infini::ccl
+
+#endif  // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_GATHER_H_

src/backends/ccl/nccl/api.h

@@ -47,6 +47,13 @@ struct NcclApi {
                          stream);
   }
 
+  static Result AllGather(const void* send_buff, void* recv_buff,
+                          size_t send_count, DataType data_type, Comm comm,
+                          Stream stream) {
+    return ncclAllGather(send_buff, recv_buff, send_count, data_type, comm,
+                         stream);
+  }
+
   static Result Send(const void* send_buff, size_t count, DataType data_type,
                      int peer, Comm comm, Stream stream) {
     return ncclSend(send_buff, count, data_type, peer, comm, stream);

src/backends/ccl/nccl/impl/all_gather.h

@@ -0,0 +1,17 @@
+#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_GATHER_H_
+#define INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_GATHER_H_
+
+#include "backends/ccl/common/impl/all_gather.h"
+
+namespace infini::ccl {
+
+template <Device::Type device>
+class AllGatherImpl<BackendType::kNccl, device>
+    : public CclAllGatherImpl<BackendType::kNccl, device> {};
+
+template <>
+struct BackendEnabled<AllGather, BackendType::kNccl> : std::true_type {};
+
+}  // namespace infini::ccl
+
+#endif  // INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_GATHER_H_

src/base/all_gather.h

@@ -19,30 +19,65 @@ class AllGather : public Operation<AllGather> {
   static ReturnStatus Execute(const void *send_buff, void *recv_buff,
                               size_t count, DataType datatype,
                               void *comm_handle, void *stream) {
-    if (HasInvalidArgs(send_buff, recv_buff, datatype, comm_handle)) {
+    if (!comm_handle) {
+      LOG("Invalid communicator handle for `AllGather`.");
       return ReturnStatus::kInvalidArgument;
     }
+
     auto *comm = static_cast<Communicator *>(comm_handle);
+    if (HasInvalidArgs(send_buff, recv_buff, count, datatype)) {
+      return ReturnStatus::kInvalidArgument;
+    }
+    if (count == 0) {
+      return ReturnStatus::kSuccess;
+    }
+
+    if (!comm->HasBackend(backend_type) || comm->device_type() != device_type) {
+      using DispatchKey =
+          typename BackendDependentType<backend_type, AllGather>::type;
+      const BackendType comm_backend =
+          Operation<DispatchKey>::FindSupportedBackend(
+              comm->device_type(),
+              {comm->HasBackend(backend_type) ? backend_type
+                                              : BackendType::kCount,
+               comm->intra_comm_backend(), comm->inter_comm_backend()});
+      if (comm_backend == BackendType::kCount) {
+        if (comm->intra_comm_backend() == BackendType::kCount &&
+            comm->inter_comm_backend() == BackendType::kCount) {
+          LOG("No initialized backend is available for `AllGather`.");
+          return ReturnStatus::kInternalError;
+        }
+        return ReturnStatus::kNotSupported;
+      }
+
+      return Operation<DispatchKey>::Call(comm_backend, comm->device_type(),
+                                          send_buff, recv_buff, count, datatype,
+                                          comm_handle, stream);
+    }
+
     return AllGatherImpl<backend_type, device_type>::Apply(
         send_buff, recv_buff, count, datatype, comm, stream);
   }
 
  private:
+  template <BackendType, typename T>
+  struct BackendDependentType {
+    using type = T;
+  };
+
   static bool HasInvalidArgs(const void *send_buff, void *recv_buff,
-                             DataType datatype, void *comm_handle) {
-    if (!comm_handle) {
-      // TODO(lzm): change to use `glog`.
-      LOG("Invalid communicator handle for `AllGather`.");
+                             size_t count, DataType datatype) {
+    if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) {
+      LOG("Invalid data type for `AllGather`.");
       return true;
     }
+    if (count == 0) {
+      return false;
+    }
     if (!send_buff || !recv_buff) {
       LOG("Invalid buffer pointer for `AllGather`.");
       return true;
     }
-    if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) {
-      LOG("Invalid data type for `AllGather`.");
-      return true;
-    }
     return false;
   }
 };