feat: support NCCL AllReduce

.github/pull_request_template.md

@@ -63,6 +63,7 @@ Please check all the platforms and/or backends this PR affects (i.e., code is to
 
 - [ ] OpenMPI
 - [ ] MPICH
+- [ ] NCCL
 
 ## Performance Impact
 
@@ -110,6 +111,7 @@ See `CONTRIBUTING.md` § Pull Requests for the official testing requirements and
 
 - [ ] OpenMPI
 - [ ] MPICH
+- [ ] NCCL
 
 ---
 

README.md

@@ -352,6 +352,7 @@ export LD_LIBRARY_PATH=${INFINI_INSTALL}/lib:$LD_LIBRARY_PATH
 |---------|---------------|----------------------|---------------|
 | **OpenMPI** | Full | `WITH_OMPI=ON` | The default backend. Requires the OpenMPI development package.|
 | **MPICH** | Full | `WITH_MPICH=ON` | Requires the MPICH development package.|
+| **NCCL** | Partial | `WITH_NCCL=ON` | Requires NVIDIA NCCL. Currently available only when `WITH_NVIDIA=ON`.|
 
 </details>
 

examples/ccl/all_reduce.cc

@@ -0,0 +1,190 @@
+/**
+ * InfiniCCL Example: Thread-per-GPU Single-Node AllReduce
+ *
+ * This example demonstrates spawning one CPU thread per GPU on a single node,
+ * generating a shared UniqueID, and performing an AllReduce entirely via
+ * InfiniCCL's native CCL backend without any OpenMPI dependencies.
+ *
+ * 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(infinicclAllReduce(d_send, d_recv, args.num_elements,
+                                    infinicclFloat32, infinicclSum, comm,
+                                    nullptr));
+  }
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+
+  // Profiling Iterations
+  Timer timer;
+  for (int i = 0; i < args.profile_iter; i++) {
+    CHECK_INFINI(infinicclAllReduce(d_send, d_recv, args.num_elements,
+                                    infinicclFloat32, infinicclSum, comm,
+                                    nullptr));
+  }
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+  double elapsed = timer.ElapsedMs() / static_cast<double>(args.profile_iter);
+
+  // Copy output reduced data back from device to host.
+  CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes,
+                          Rt::MemcpyDeviceToHost));
+
+  // Result Validation
+  float expected = 0.0f;
+  for (int r = 0; r < args.size; r++) {
+    expected += static_cast<float>(r + 1);
+  }
+  Validator::ValidateResult(h_recv.data(), args.num_elements, expected,
+                            args.rank, true, "AllReduce");
+
+  // Metrics Reporting (Only Rank 0)
+  if (args.rank == 0) {
+    std::cout << "\n=== Single-Node Threaded AllReduce 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;
+
+  int opt;
+  while ((opt = getopt(argc, argv, "g:w:p:n:h")) != -1) {
+    switch (opt) {
+      case 'g':
+        num_gpus = std::stoi(optarg);
+        break;
+      case 'w':
+        warmup_iters = std::stoi(optarg);
+        break;
+      case 'p':
+        profile_iters = std::stoi(optarg);
+        break;
+      case 'n':
+        num_elements = static_cast<size_t>(std::stoull(optarg));
+        break;
+      case 'h':
+        std::cout << "Usage: " << argv[0] << " [options]\n"
+                  << "Options:\n"
+                  << "  -g <num_gpus>        Number of GPUs (default: 8)\n"
+                  << "  -w <warmup_iters>    Warmup iterations (default: 1)\n"
+                  << "  -p <profile_iters>   Profile iterations (default: 20)\n"
+                  << "  -n <num_elements>    Number of elements (default: "
+                  << (1 << 25) << ")\n";
+        return EXIT_SUCCESS;
+      default:
+        std::cerr << "Invalid argument. Use -h for help." << std::endl;
+        return EXIT_FAILURE;
+    }
+  }
+
+  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;
+}

examples/ccl_mpi_hybrid/all_reduce.cc

@@ -0,0 +1,161 @@
+/**
+ * InfiniCCL Example: AllReduce (Ompi + CCL Hybrid)
+ * * This example demonstrates a program that performs a
+ * collective sum-reduction across multiple GPUs and nodes.
+ */
+
+#include <unistd.h>
+
+#include <iostream>
+#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"
+#include "device.h"
+#include "runtime.h"
+#include "traits.h"
+
+using namespace infini::ccl;
+
+void RunAllReduceExample(int argc, char **argv, int warmup_iter,
+                         int profile_iter, const size_t kNumElements) {
+  constexpr Device::Type kDevType =
+      ListGetBest<DevicePriority>(EnabledDevices{});
+  using Rt = Runtime<kDevType>;
+
+  CHECK_INFINI(infinicclInit(&argc, &argv));
+
+  int rank, size;
+  CHECK_INFINI(infinicclGetRank(&rank));
+  CHECK_INFINI(infinicclGetSize(&size));
+
+  char hostname[256];
+  gethostname(hostname, sizeof(hostname));
+
+  // Map local rank to GPU device.
+  // Note: this is just for info printing. In practice, this part is not needed.
+  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 Communicator
+  infinicclUniqueId id;
+  infinicclComm_t comm = nullptr;
+  CHECK_INFINI(infinicclCommInitAll(&comm, size, nullptr));
+
+  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: " << Device::StringFromType(kDevType) << " "
+            << " | Device " << local_rank << std::endl;
+
+  CHECK_INFINI(infinicclCommInitRank(&comm, size, id, rank));
+
+  // Prepare Data
+  std::vector<float> h_send(kNumElements);
+  std::vector<float> h_recv(kNumElements, 0.0f);
+
+  // Initialize: each rank provides its (rank + 1) as data.
+  for (size_t i = 0; i < kNumElements; i++) {
+    h_send[i] = static_cast<float>(rank + 1);
+  }
+
+  float *d_send, *d_recv;
+  size_t total_bytes = kNumElements * sizeof(*d_send);
+  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));
+
+  if (rank == 0) {
+    std::cout << "\n=== Performing AllReduce on GPU Memory ===" << std::endl;
+    std::cout << "Data size: " << kNumElements << " floats ("
+              << total_bytes / 1024 / 1024 << " MB)" << std::endl;
+    std::cout << "Operation: Sum" << 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));
+
+  // Warm-up and D2H transfer the answer.
+  CHECK_INFINI(infinicclAllReduce(d_send, d_recv, kNumElements,
+                                  infinicclFloat32, infinicclSum, comm,
+                                  nullptr));
+  CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, kNumElements * sizeof(float),
+                          Rt::MemcpyDeviceToHost));
+
+  for (int i = 1; i < warmup_iter; ++i) {
+    CHECK_INFINI(infinicclAllReduce(d_send, d_recv, kNumElements,
+                                    infinicclFloat32, infinicclSum, comm,
+                                    nullptr));
+  }
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+
+  // Profiling
+  Timer timer;
+
+  for (int i = 0; i < profile_iter; i++) {
+    CHECK_INFINI(infinicclAllReduce(d_send, d_recv, kNumElements,
+                                    infinicclFloat32, infinicclSum, comm,
+                                    nullptr));
+  }
+
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+  double elapsed = timer.ElapsedMs() / static_cast<double>(profile_iter);
+
+  // Result Validation
+  float expected = 0.0f;
+  for (int r = 0; r < size; r++) {
+    expected += static_cast<float>(r + 1);
+  }
+
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+  Validator::ValidateResult(h_recv.data(), kNumElements, expected, rank, true,
+                            "AllReduce");
+
+  // Metrics Reporting (Only from rank 0 for cleaner output)
+  if (rank == 0) {
+    Metrics metrics{elapsed, total_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;
+  }
+}
+
+int main(int argc, char **argv) {
+  int warmup_iters = 2;
+  int profile_iters = 20;
+  size_t num_elements = 1 << 25;
+
+  RunAllReduceExample(argc, argv, warmup_iters, profile_iters, num_elements);
+
+  return EXIT_SUCCESS;
+}