From 3816cd5d93dcfa0332430e7afa3e323bca1bbf84 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 12 Aug 2026 05:20:15 +0800 Subject: [PATCH 1/3] feat(nccl): add all-gather communication --- examples/ccl/all_gather.cc | 123 ++++++++++++++++++++++ examples/mpi/all_gather.cc | 17 +-- src/backends/ccl/common/impl/all_gather.h | 44 ++++++++ src/backends/ccl/nccl/api.h | 7 ++ src/backends/ccl/nccl/impl/all_gather.h | 17 +++ src/base/all_gather.h | 53 ++++++++-- 6 files changed, 244 insertions(+), 17 deletions(-) create mode 100644 examples/ccl/all_gather.cc create mode 100644 src/backends/ccl/common/impl/all_gather.h create mode 100644 src/backends/ccl/nccl/impl/all_gather.h diff --git a/examples/ccl/all_gather.cc b/examples/ccl/all_gather.cc new file mode 100644 index 0000000..e411172 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include + +#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& output, int rank, const char* mode) { + for (int source = 0; source < kRankCount; ++source) { + const float expected = static_cast(source + 1); + const size_t offset = static_cast(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(EnabledDevices{}); + using Rt = Runtime; + + CHECK_RT(Rt, Rt::SetDevice(args.rank)); + + infinicclComm_t comm = nullptr; + CHECK_INFINI(infinicclCommInitRank(&comm, kRankCount, args.id, args.rank)); + + std::vector host_send(kNumElements, static_cast(args.rank + 1)); + std::vector 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(&device_send), send_bytes)); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&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(args.rank) * kNumElements, + kNumElements, static_cast(args.rank + 1)); + CHECK_RT(Rt, Rt::Memcpy(device_recv, host_recv.data(), recv_bytes, + Rt::MemcpyHostToDevice)); + + float* local_block = + device_recv + static_cast(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 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; +} diff --git a/examples/mpi/all_gather.cc b/examples/mpi/all_gather.cc index c0f1a62..5716b20 100644 --- a/examples/mpi/all_gather.cc +++ b/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(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(src_rank + 1); size_t offset = static_cast(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,6 +158,7 @@ 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) { @@ -166,7 +166,8 @@ int main(int argc, char **argv) { 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; } diff --git a/src/backends/ccl/common/impl/all_gather.h b/src/backends/ccl/common/impl/all_gather.h new file mode 100644 index 0000000..ea8c95c --- /dev/null +++ b/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 +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; + using TypeMap = CclTypeMap; + using CommInstance = CclCommInstance; + + if (!comm || !comm->intra_comm() || comm->intra_comm_backend() != backend || + comm->device_type() != device) { + return ReturnStatus::kInternalError; + } + + auto* intra = static_cast(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(stream))); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_GATHER_H_ diff --git a/src/backends/ccl/nccl/api.h b/src/backends/ccl/nccl/api.h index ed6e62f..e1d51a3 100644 --- a/src/backends/ccl/nccl/api.h +++ b/src/backends/ccl/nccl/api.h @@ -49,6 +49,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); diff --git a/src/backends/ccl/nccl/impl/all_gather.h b/src/backends/ccl/nccl/impl/all_gather.h new file mode 100644 index 0000000..fdabdd7 --- /dev/null +++ b/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 +class AllGatherImpl + : public CclAllGatherImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_GATHER_H_ diff --git a/src/base/all_gather.h b/src/base/all_gather.h index c0072d8..c7ce6ff 100644 --- a/src/base/all_gather.h +++ b/src/base/all_gather.h @@ -19,30 +19,65 @@ class AllGather : public Operation { 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(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::type; + const BackendType comm_backend = + Operation::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::Call(comm_backend, comm->device_type(), + send_buff, recv_buff, count, datatype, + comm_handle, stream); + } + return AllGatherImpl::Apply( send_buff, recv_buff, count, datatype, comm, stream); } private: + template + 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; } }; From c61349b2fb767aab7efc551c45da851ac573b05e Mon Sep 17 00:00:00 2001 From: Zimin Li Date: Tue, 15 Sep 2026 08:14:30 +0000 Subject: [PATCH 2/3] fix: simplify all-gather dispatch and add its cross-node hybrid example program - remove communicator redispatch from the `AllGather` operation - align CCL `AllGather` communicator validation with the existing collective style - handle zero-element operations without requiring data buffers - add a hybrid example for global-rank `AllGather` through InfiniCCL APIs --- examples/ccl_mpi_hybrid/all_gather.cc | 183 ++++++++++++++++++++++ src/backends/ccl/common/impl/all_gather.h | 5 +- src/base/all_gather.h | 53 ++----- 3 files changed, 197 insertions(+), 44 deletions(-) create mode 100644 examples/ccl_mpi_hybrid/all_gather.cc diff --git a/examples/ccl_mpi_hybrid/all_gather.cc b/examples/ccl_mpi_hybrid/all_gather.cc new file mode 100644 index 0000000..aa2d5be --- /dev/null +++ b/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 + +#include +#include +#include +#include +#include + +#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::EnabledDevices{}); + using Rt = ccl::Runtime; + + // 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 h_send(num_elements, static_cast(rank + 1)); + std::vector h_recv(static_cast(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); + + 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(profile_iter); + + // Result Validation + bool correct = true; + for (int source_rank = 0; source_rank < size; ++source_rank) { + const float expected = static_cast(source_rank + 1); + const size_t offset = static_cast(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(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; +} diff --git a/src/backends/ccl/common/impl/all_gather.h b/src/backends/ccl/common/impl/all_gather.h index ea8c95c..10fc340 100644 --- a/src/backends/ccl/common/impl/all_gather.h +++ b/src/backends/ccl/common/impl/all_gather.h @@ -18,13 +18,12 @@ class CclAllGatherImpl { using TypeMap = CclTypeMap; using CommInstance = CclCommInstance; - if (!comm || !comm->intra_comm() || comm->intra_comm_backend() != backend || - comm->device_type() != device) { + if (!comm) { return ReturnStatus::kInternalError; } auto* intra = static_cast(comm->intra_comm()); - if (!intra->handle) { + if (!intra || !intra->handle) { return ReturnStatus::kInternalError; } diff --git a/src/base/all_gather.h b/src/base/all_gather.h index c7ce6ff..1341cf4 100644 --- a/src/base/all_gather.h +++ b/src/base/all_gather.h @@ -19,65 +19,36 @@ class AllGather : public Operation { static ReturnStatus Execute(const void *send_buff, void *recv_buff, size_t count, DataType datatype, void *comm_handle, void *stream) { - if (!comm_handle) { - LOG("Invalid communicator handle for `AllGather`."); - return ReturnStatus::kInvalidArgument; - } - - auto *comm = static_cast(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::type; - const BackendType comm_backend = - Operation::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::Call(comm_backend, comm->device_type(), - send_buff, recv_buff, count, datatype, - comm_handle, stream); + if (HasInvalidArgs(send_buff, recv_buff, datatype, comm_handle)) { + return ReturnStatus::kInvalidArgument; } + auto *comm = static_cast(comm_handle); + return AllGatherImpl::Apply( send_buff, recv_buff, count, datatype, comm, stream); } private: - template - struct BackendDependentType { - using type = T; - }; - static bool HasInvalidArgs(const void *send_buff, void *recv_buff, - size_t count, DataType datatype) { - if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) { - LOG("Invalid data type for `AllGather`."); + DataType datatype, void *comm_handle) { + if (!comm_handle) { + // TODO(lzm): change to use `glog`. + LOG("Invalid communicator handle 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; } }; From abc004bc08e198d9c2052b2378cb375d0cd157e2 Mon Sep 17 00:00:00 2001 From: Zimin Li Date: Tue, 15 Sep 2026 10:08:41 +0000 Subject: [PATCH 3/3] feat: improve all-gather validation and benchmarking - add configurable element count, warmup iterations, and profiling iterations - print clear pass/fail summaries for out-of-place and in-place all-gather modes --- examples/ccl/all_gather.cc | 123 +++++++++++++++++++++++++++++-------- 1 file changed, 97 insertions(+), 26 deletions(-) diff --git a/examples/ccl/all_gather.cc b/examples/ccl/all_gather.cc index e411172..9f640d7 100644 --- a/examples/ccl/all_gather.cc +++ b/examples/ccl/all_gather.cc @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -22,29 +23,49 @@ using namespace infini::ccl; namespace { constexpr int kRankCount = 2; -constexpr size_t kNumElements = 1 << 10; 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& output, int rank, const char* mode) { +bool Validate(const std::vector& 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(source + 1); - const size_t offset = static_cast(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; - } - } + const size_t offset = static_cast(source) * num_elements; + + correct = Validator::ValidateResult(output.data() + offset, num_elements, + expected, rank) && + correct; + } + return correct; +} + +void PrintResult(const std::vector& 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(source) * num_elements; + std::cout << "[r" << source << ": " << output[offset] << "] "; } - return true; + std::cout << std::endl; } void WorkerThread(ThreadArgs args) { @@ -57,42 +78,80 @@ void WorkerThread(ThreadArgs args) { infinicclComm_t comm = nullptr; CHECK_INFINI(infinicclCommInitRank(&comm, kRankCount, args.id, args.rank)); - std::vector host_send(kNumElements, static_cast(args.rank + 1)); - std::vector host_recv(kNumElements * kRankCount, 0.0f); + std::vector host_send(args.num_elements, + static_cast(args.rank + 1)); + std::vector 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 = kNumElements * sizeof(float); + 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(&device_send), send_bytes)); CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&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(args.profile_iter); - 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")) { + + 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(args.rank) * kNumElements, - kNumElements, static_cast(args.rank + 1)); + std::fill_n( + host_recv.begin() + static_cast(args.rank) * args.num_elements, + args.num_elements, static_cast(args.rank + 1)); CHECK_RT(Rt, Rt::Memcpy(device_recv, host_recv.data(), recv_bytes, Rt::MemcpyHostToDevice)); float* local_block = - device_recv + static_cast(args.rank) * kNumElements; - CHECK_INFINI(infinicclAllGather(local_block, device_recv, kNumElements, + device_recv + static_cast(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)); - if (!Validate(host_recv, args.rank, "In-place AllGather")) { + 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)); @@ -101,23 +160,35 @@ void WorkerThread(ThreadArgs args) { } // 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 threads; threads.reserve(kRankCount); + for (int rank = 0; rank < kRankCount; ++rank) { - threads.emplace_back(WorkerThread, - ThreadArgs{rank, shared_id, &all_correct}); + 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; }