feat(nccl): add all-gather communication

examples/ccl/all_gather.cc

@@ -0,0 +1,194 @@
+/**
+ * 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 <string>
+#include <thread>
+#include <vector>
+
+#include "backend_manifest.h"
+#include "infiniccl.h"
+#include "utils.h"
+
+using namespace infini::ccl;
+
+namespace {
+
+constexpr int kRankCount = 2;
+
+struct ThreadArgs {
+  int rank;
+  infinicclUniqueId id;
+  size_t num_elements;
+  int warmup_iter;
+  int profile_iter;
+  std::atomic_bool* all_correct;
+};
+
+bool Validate(const std::vector<float>& output, int rank) {
+  // Every rank checks the complete gathered output by source-rank block.
+  const size_t num_elements = output.size() / kRankCount;
+  bool correct = true;
+  for (int source = 0; source < kRankCount; ++source) {
+    const float expected = static_cast<float>(source + 1);
+    const size_t offset = static_cast<size_t>(source) * num_elements;
+
+    correct = Validator::ValidateResult(output.data() + offset, num_elements,
+                                        expected, rank) &&
+              correct;
+  }
+  return correct;
+}
+
+void PrintResult(const std::vector<float>& output, const char* mode,
+                 bool correct) {
+  const char* green = "\033[32m";
+  const char* red = "\033[31m";
+  const char* reset = "\033[0m";
+  const size_t num_elements = output.size() / kRankCount;
+
+  std::cout << "\n=== " << mode << " AllGather Results ===" << std::endl;
+  std::cout << "Correct: "
+            << (correct ? (green + std::string("YES") + reset)
+                        : (red + std::string("NO") + reset))
+            << std::endl;
+  std::cout << "Sample blocks: ";
+  for (int source = 0; source < kRankCount; ++source) {
+    const size_t offset = static_cast<size_t>(source) * num_elements;
+    std::cout << "[r" << source << ": " << output[offset] << "] ";
+  }
+  std::cout << std::endl;
+}
+
+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(args.num_elements,
+                               static_cast<float>(args.rank + 1));
+  std::vector<float> host_recv(args.num_elements * kRankCount, 0.0f);
+
+  // Prepare separate send and receive buffers for the out-of-place case.
+  float* device_send = nullptr;
+  float* device_recv = nullptr;
+  const size_t send_bytes = args.num_elements * 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_RT(Rt, Rt::StreamSynchronize(nullptr));
+
+  auto all_gather = [&]() {
+    return infinicclAllGather(device_send, device_recv, args.num_elements,
+                              infinicclFloat32, comm, nullptr);
+  };
+
+  for (int i = 0; i < args.warmup_iter; ++i) {
+    CHECK_INFINI(all_gather());
+  }
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+
+  Timer timer;
+  for (int i = 0; i < args.profile_iter; ++i) {
+    CHECK_INFINI(all_gather());
+  }
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+  const double elapsed =
+      timer.ElapsedMs() / static_cast<double>(args.profile_iter);
+
+  CHECK_RT(Rt, Rt::Memcpy(host_recv.data(), device_recv, recv_bytes,
+                          Rt::MemcpyDeviceToHost));
+
+  const bool out_of_place_correct = Validate(host_recv, args.rank);
+  if (!out_of_place_correct) {
+    args.all_correct->store(false, std::memory_order_relaxed);
+  }
+  if (args.rank == 0) {
+    PrintResult(host_recv, "Out-of-place", out_of_place_correct);
+  }
+
+  // Seed only the local block, then use it as both input and output.
+  std::fill(host_recv.begin(), host_recv.end(), 0.0f);
+  std::fill_n(
+      host_recv.begin() + static_cast<size_t>(args.rank) * args.num_elements,
+      args.num_elements, 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) * args.num_elements;
+  CHECK_INFINI(infinicclAllGather(local_block, device_recv, args.num_elements,
+                                  infinicclFloat32, comm, nullptr));
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+  CHECK_RT(Rt, Rt::Memcpy(host_recv.data(), device_recv, recv_bytes,
+                          Rt::MemcpyDeviceToHost));
+  const bool in_place_correct = Validate(host_recv, args.rank);
+  if (!in_place_correct) {
+    args.all_correct->store(false, std::memory_order_relaxed);
+  }
+
+  if (args.rank == 0) {
+    PrintResult(host_recv, "In-place", in_place_correct);
+    std::cout << "\n=== Single-Node Threaded AllGather Results ==="
+              << std::endl;
+    Metrics metrics{elapsed, recv_bytes, kRankCount};
+    metrics.Print();
+  }
+
+  // Cleanup rank-local resources.
+  CHECK_RT(Rt, Rt::Free(device_send));
+  CHECK_RT(Rt, Rt::Free(device_recv));
+  CHECK_INFINI(infinicclCommDestroy(comm));
+}
+
+}  // namespace
+
+int main() {
+  constexpr size_t kNumElements = 1 << 20;
+  constexpr int kWarmupIterations = 2;
+  constexpr int kProfileIterations = 20;
+
+  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) {
+    ThreadArgs args{rank,
+                    shared_id,
+                    kNumElements,
+                    kWarmupIterations,
+                    kProfileIterations,
+                    &all_correct};
+    threads.emplace_back(WorkerThread, args);
+  }
+  for (auto& thread : threads) {
+    thread.join();
+  }
+
+  if (!all_correct.load(std::memory_order_relaxed)) {
+    std::cerr << "AllGather validation failed." << std::endl;
+    return EXIT_FAILURE;
+  }
+
+  std::cout << "AllGather validation passed." << std::endl;
+  return EXIT_SUCCESS;
+}

examples/ccl_mpi_hybrid/all_gather.cc

@@ -0,0 +1,183 @@
+/**
+ * InfiniCCL Example: AllGather (MPI + CCL Hybrid)
+ *
+ * This example performs a collective gather across global ranks. MPI
+ * bootstraps the processes and the CCL unique ID; the collective itself uses
+ * the CCL backend.
+ */
+
+#include <unistd.h>
+
+#include <algorithm>
+#include <cstdlib>
+#include <iostream>
+#include <string>
+#include <vector>
+
+#include "backend_manifest.h"
+#include "device.h"
+#include "infiniccl.h"
+#include "runtime.h"
+#include "traits.h"
+#include "utils.h"
+
+namespace ccl = infini::ccl;
+
+bool RunAllGatherExample(int argc, char **argv, int warmup_iter,
+                         int profile_iter, size_t num_elements) {
+  constexpr ccl::Device::Type kDevType =
+      ccl::ListGetBest<ccl::DevicePriority>(ccl::EnabledDevices{});
+  using Rt = ccl::Runtime<kDevType>;
+
+  // Initialize InfiniCCL and obtain the global rank information.
+  CHECK_INFINI(infinicclInit(&argc, &argv));
+
+  int rank = 0;
+  int size = 0;
+  CHECK_INFINI(infinicclGetRank(&rank));
+  CHECK_INFINI(infinicclGetSize(&size));
+
+  char hostname[256];
+  gethostname(hostname, sizeof(hostname));
+
+  // Map local rank to GPU device.
+  const char *local_rank_str = std::getenv("OMPI_COMM_WORLD_LOCAL_RANK");
+  int local_rank = 0;
+  if (local_rank_str != nullptr) {
+    local_rank = std::atoi(local_rank_str);
+  }
+
+  CHECK_RT(Rt, Rt::SetDevice(local_rank));
+
+  // Setup the MPI-backed communicator used to bootstrap the CCL unique ID.
+  infinicclComm_t comm = nullptr;
+  CHECK_INFINI(infinicclCommInitAll(&comm, size, nullptr));
+
+  infinicclUniqueId id;
+  if (rank == 0) {
+    CHECK_INFINI(infinicclGetUniqueId(&id));
+  }
+  CHECK_INFINI(infinicclBroadcast(&id, &id, sizeof(id), infinicclChar, 0, comm,
+                                  nullptr));
+
+  std::cout << "[Rank " << rank << "] Host: " << hostname
+            << " | GPU: " << ccl::Device::StringFromType(kDevType) << " "
+            << " | Device " << local_rank << std::endl;
+
+  CHECK_INFINI(infinicclCommInitRank(&comm, size, id, rank));
+
+  // Prepare Data
+  std::vector<float> h_send(num_elements, static_cast<float>(rank + 1));
+  std::vector<float> h_recv(static_cast<size_t>(size) * num_elements, 0.0f);
+
+  float *d_send = nullptr;
+  float *d_recv = nullptr;
+  const size_t send_bytes = num_elements * sizeof(*d_send);
+  const size_t recv_bytes = send_bytes * static_cast<size_t>(size);
+
+  CHECK_RT(Rt, Rt::Malloc((void **)&d_send, send_bytes));
+  CHECK_RT(Rt, Rt::Malloc((void **)&d_recv, recv_bytes));
+  CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), send_bytes,
+                          Rt::MemcpyHostToDevice));
+  CHECK_RT(Rt, Rt::Memcpy(d_recv, h_recv.data(), recv_bytes,
+                          Rt::MemcpyHostToDevice));
+
+  if (rank == 0) {
+    std::cout << "\n=== Performing AllGather on GPU Memory ===" << std::endl;
+    std::cout << "Data size: " << num_elements << " floats ("
+              << send_bytes / 1024 / 1024 << " MB per rank)" << std::endl;
+    std::cout << "Operation: AllGather" << std::endl;
+    std::cout << "Warm-up iterations: " << warmup_iter << std::endl;
+    std::cout << "Profile iterations: " << profile_iter << std::endl;
+  }
+
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+
+  auto all_gather_call = [&]() {
+    return infinicclAllGather(d_send, d_recv, num_elements, infinicclFloat32,
+                              comm, nullptr);
+  };
+
+  // Warm-up and D2H transfer the answer.
+  CHECK_INFINI(all_gather_call());
+  CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, recv_bytes,
+                          Rt::MemcpyDeviceToHost));
+
+  for (int i = 1; i < warmup_iter; ++i) {
+    CHECK_INFINI(all_gather_call());
+  }
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+
+  // Profiling
+  Timer timer;
+
+  for (int i = 0; i < profile_iter; ++i) {
+    CHECK_INFINI(all_gather_call());
+  }
+
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+  CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, recv_bytes,
+                          Rt::MemcpyDeviceToHost));
+
+  const double elapsed = timer.ElapsedMs() / static_cast<double>(profile_iter);
+
+  // Result Validation
+  bool correct = true;
+  for (int source_rank = 0; source_rank < size; ++source_rank) {
+    const float expected = static_cast<float>(source_rank + 1);
+    const size_t offset = static_cast<size_t>(source_rank) * num_elements;
+
+    correct = Validator::ValidateResult(h_recv.data() + offset, num_elements,
+                                        expected, rank) &&
+              correct;
+  }
+
+  if (rank == 0) {
+    const char *kGreen = "\033[32m";
+    const char *kRed = "\033[31m";
+    const char *kReset = "\033[0m";
+
+    std::cout << "\n=== AllGather Results ===" << std::endl;
+    std::cout << "Correct: "
+              << (correct ? (kGreen + std::string("YES") + kReset)
+                          : (kRed + std::string("NO") + kReset))
+              << std::endl;
+
+    std::cout << "Sample blocks: ";
+    for (int source_rank = 0; source_rank < std::min(size, 4); ++source_rank) {
+      const size_t offset = static_cast<size_t>(source_rank) * num_elements;
+      std::cout << "[r" << source_rank << ": " << h_recv[offset] << "] ";
+    }
+    std::cout << std::endl;
+  }
+
+  // Metrics Reporting (Only from rank 0 for cleaner output)
+  if (rank == 0) {
+    Metrics metrics{elapsed, recv_bytes, size};
+    metrics.Print();
+  }
+
+  // Cleanup
+  CHECK_RT(Rt, Rt::Free(d_send));
+  CHECK_RT(Rt, Rt::Free(d_recv));
+
+  CHECK_INFINI(infinicclCommDestroy(comm));
+  CHECK_INFINI(infinicclFinalize());
+
+  if (rank == 0) {
+    std::cout << "InfiniCCL finalized." << std::endl;
+  }
+
+  return correct;
+}
+
+int main(int argc, char **argv) {
+  const int warmup_iters = 2;
+  const int profile_iters = 20;
+  const size_t num_elements = 1 << 20;
+
+  return RunAllGatherExample(argc, argv, warmup_iters, profile_iters,
+                             num_elements)
+             ? EXIT_SUCCESS
+             : EXIT_FAILURE;
+}