feat(nccl): add point-to-point communication

examples/ccl/send_recv.cc

@@ -0,0 +1,134 @@
+/**
+ * InfiniCCL Example: Thread-per-GPU Single-Node Send/Recv
+ *
+ * This example transfers data from GPU 0 to GPU 1 through InfiniCCL's native
+ * CCL backend without an MPI launcher.
+ */
+
+#include <atomic>
+#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 int kSender = 0;
+constexpr int kReceiver = 1;
+constexpr float kSendValue = 7.0f;
+
+struct ThreadArgs {
+  int rank;
+  infinicclUniqueId id;
+  size_t num_elements;
+  int warmup_iter;
+  int profile_iter;
+  std::atomic_bool* all_correct;
+};
+
+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_buffer(args.num_elements,
+                                 args.rank == kSender ? kSendValue : 0.0f);
+  float* device_buffer = nullptr;
+  const size_t total_bytes = args.num_elements * sizeof(float);
+
+  CHECK_RT(Rt,
+           Rt::Malloc(reinterpret_cast<void**>(&device_buffer), total_bytes));
+  CHECK_RT(Rt, Rt::Memcpy(device_buffer, host_buffer.data(), total_bytes,
+                          Rt::MemcpyHostToDevice));
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+
+  auto exchange = [&]() {
+    if (args.rank == kSender) {
+      return infinicclSend(device_buffer, args.num_elements, infinicclFloat32,
+                           kReceiver, comm, nullptr);
+    }
+    return infinicclRecv(device_buffer, args.num_elements, infinicclFloat32,
+                         kSender, comm, nullptr);
+  };
+
+  for (int i = 0; i < args.warmup_iter; ++i) {
+    CHECK_INFINI(exchange());
+  }
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+
+  Timer timer;
+  for (int i = 0; i < args.profile_iter; ++i) {
+    CHECK_INFINI(exchange());
+  }
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+  const double elapsed =
+      timer.ElapsedMs() / static_cast<double>(args.profile_iter);
+
+  if (args.rank == kReceiver) {
+    CHECK_RT(Rt, Rt::Memcpy(host_buffer.data(), device_buffer, total_bytes,
+                            Rt::MemcpyDeviceToHost));
+    const bool correct =
+        Validator::ValidateResult(host_buffer.data(), args.num_elements,
+                                  kSendValue, kSender, true, "Send/Recv");
+    if (!correct) {
+      args.all_correct->store(false, std::memory_order_relaxed);
+    }
+  } else {
+    std::cout << "\n=== Single-Node Threaded Send/Recv Results ==="
+              << std::endl;
+    Metrics metrics{elapsed, total_bytes, kRankCount};
+    metrics.Print();
+  }
+
+  CHECK_RT(Rt, Rt::Free(device_buffer));
+  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 << "Send/Recv validation failed." << std::endl;
+    return EXIT_FAILURE;
+  }
+
+  std::cout << "Send/Recv validation passed." << std::endl;
+  return EXIT_SUCCESS;
+}

examples/ccl_mpi_hybrid/send_recv.cc

@@ -0,0 +1,204 @@
+/**
+ * InfiniCCL Example: Send/Recv (Ompi + CCL Hybrid)
+ *
+ * This example performs point-to-point `infinicclSend` and `infinicclRecv`
+ * operations between global ranks. OpenMPI bootstraps the processes and the
+ * CCL unique ID; the data transfer itself uses the CCL backend.
+ */
+
+#include <unistd.h>
+
+#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;
+
+void RunSendRecvExample(int argc, char **argv, int warmup_iter,
+                        int profile_iter, size_t num_elements, int sender,
+                        int receiver, int required_ranks, float send_value) {
+  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));
+
+  if (size < required_ranks) {
+    if (rank == sender) {
+      std::cerr << "Send/Recv example requires at least " << required_ranks
+                << " ranks." << std::endl;
+    }
+
+    CHECK_INFINI(infinicclFinalize());
+    return;
+  }
+
+  // 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 == sender) {
+    CHECK_INFINI(infinicclGetUniqueId(&id));
+  }
+  CHECK_INFINI(infinicclBroadcast(&id, &id, sizeof(id), infinicclChar, sender,
+                                  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, send_value);
+  std::vector<float> h_recv(num_elements, 0.0f);
+
+  float *d_send = nullptr;
+  float *d_recv = nullptr;
+  size_t total_bytes = num_elements * 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 == sender) {
+    std::cout << "\n=== Performing Send/Recv on GPU Memory ===" << std::endl;
+    std::cout << "Sender rank: " << sender << std::endl;
+    std::cout << "Receiver rank: " << receiver << std::endl;
+    std::cout << "Data size: " << num_elements << " floats ("
+              << total_bytes / 1024 / 1024 << " MB)" << 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 send_recv_call = [&]() {
+    if (rank == sender) {
+      return infinicclSend(d_send, num_elements, infinicclFloat32, receiver,
+                           comm, nullptr);
+    }
+
+    if (rank == receiver) {
+      return infinicclRecv(d_recv, num_elements, infinicclFloat32, sender, comm,
+                           nullptr);
+    }
+
+    return infinicclSuccess;
+  };
+
+  // Warm-up and validate the first transfer.
+  CHECK_INFINI(send_recv_call());
+  if (rank == receiver) {
+    CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes,
+                            Rt::MemcpyDeviceToHost));
+  }
+
+  for (int i = 1; i < warmup_iter; ++i) {
+    CHECK_INFINI(send_recv_call());
+  }
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+
+  // Profiling
+  Timer timer;
+
+  for (int i = 0; i < profile_iter; i++) {
+    CHECK_INFINI(send_recv_call());
+  }
+
+  CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
+  if (rank == receiver) {
+    CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes,
+                            Rt::MemcpyDeviceToHost));
+  }
+
+  double elapsed = timer.ElapsedMs() / static_cast<double>(profile_iter);
+
+  // Result Validation
+  if (rank == receiver) {
+    bool correct = Validator::ValidateResult(
+        h_recv.data(), num_elements, send_value, rank, false, "SendRecv");
+
+    const char *kGreen = "\033[32m";
+    const char *kRed = "\033[31m";
+    const char *kReset = "\033[0m";
+
+    std::cout << "\n=== Send/Recv Results ===" << std::endl;
+    std::cout << "Correct: "
+              << (correct ? (kGreen + std::string("YES") + kReset)
+                          : (kRed + std::string("NO") + kReset))
+              << std::endl;
+    std::cout << "Expect:  " << send_value << std::endl;
+    std::cout << "Actual:  " << h_recv[0] << std::endl;
+
+    if (!correct) {
+      CHECK_RT(Rt, Rt::Free(d_send));
+      CHECK_RT(Rt, Rt::Free(d_recv));
+      CHECK_INFINI(infinicclCommDestroy(comm));
+      CHECK_INFINI(infinicclFinalize());
+      std::exit(EXIT_FAILURE);
+    }
+  }
+
+  // Metrics Reporting (Only from the sender for cleaner output)
+  if (rank == sender) {
+    Metrics metrics{elapsed, total_bytes, required_ranks};
+    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 == sender) {
+    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 << 20;
+  constexpr int kSender = 0;
+  constexpr int kReceiver = 1;
+  constexpr int kRequiredRanks = 2;
+  constexpr float kSendValue = 7.0f;
+
+  RunSendRecvExample(argc, argv, warmup_iters, profile_iters, num_elements,
+                     kSender, kReceiver, kRequiredRanks, kSendValue);
+
+  return EXIT_SUCCESS;
+}

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

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