From c1b59176db0c3967f95025592f6fc4e3d80825a7 Mon Sep 17 00:00:00 2001 From: GordonYang1 <1468121796@qq.com> Date: Mon, 17 Aug 2026 11:00:51 +0800 Subject: [PATCH] feat: support CCL `AllGather` --- examples/ccl/all_gather.cc | 309 ++++++++++++++++++++++ examples/ccl_mpi_hybrid/all_gather.cc | 293 ++++++++++++++++++++ examples/mpi/all_gather.cc | 18 +- src/backends/ccl/common/impl/all_gather.h | 63 +++++ src/backends/ccl/mccl/api.h | 7 + src/backends/ccl/mccl/impl/all_gather.h | 17 ++ src/backends/ccl/nccl/api.h | 7 + src/backends/ccl/nccl/impl/all_gather.h | 17 ++ src/backends/mpi/ompi/impl/all_gather.h | 47 +++- src/base/all_gather.h | 20 +- 10 files changed, 772 insertions(+), 26 deletions(-) create mode 100644 examples/ccl/all_gather.cc create mode 100644 examples/ccl_mpi_hybrid/all_gather.cc create mode 100644 src/backends/ccl/common/impl/all_gather.h create mode 100644 src/backends/ccl/mccl/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..44db9be --- /dev/null +++ b/examples/ccl/all_gather.cc @@ -0,0 +1,309 @@ +/** + * InfiniCCL Example: Thread-per-GPU Single-Node AllGather + * + * This example creates one native CCL rank per GPU, then validates + * out-of-place and in-place AllGather without an MPI launcher. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// 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; + +namespace { + +struct ScenarioState { + std::atomic correct{true}; + std::atomic completed{0}; +}; + +struct ThreadArgs { + int rank; + int size; + infinicclUniqueId id; + size_t num_elements; + int warmup_iter; + int profile_iter; + ScenarioState *out_of_place; + ScenarioState *in_place; +}; + +template +bool ParsePositiveNumber(const char *text, T *value) { + if (!text || !value) { + return false; + } + + T parsed{}; + const char *end = text + std::strlen(text); + const auto result = std::from_chars(text, end, parsed); + if (result.ec != std::errc{} || result.ptr != end || parsed <= 0) { + return false; + } + + *value = parsed; + return true; +} + +void PrintResult(const char *scenario, bool correct, + const std::vector &result, size_t num_elements, + int world_size, double elapsed_ms) { + constexpr const char *kGreen = "\033[32m"; + constexpr const char *kRed = "\033[31m"; + constexpr const char *kReset = "\033[0m"; + + std::cout << "\n=== " << scenario + << " CCL AllGather Results ===" << std::endl; + std::cout << "Correct: " + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) + << std::endl; + std::cout << "Average time: " << elapsed_ms << " ms" << std::endl; + std::cout << "Sample blocks: "; + for (int src_rank = 0; src_rank < world_size && src_rank < 4; ++src_rank) { + const size_t offset = static_cast(src_rank) * num_elements; + std::cout << "[r" << src_rank << ": " << result[offset] << "] "; + } + std::cout << std::endl; +} + +void WaitForScenario(ScenarioState *state, int world_size) { + state->completed.fetch_add(1, std::memory_order_release); + while (state->completed.load(std::memory_order_acquire) < world_size) { + std::this_thread::yield(); + } +} + +bool ValidateAllGather(const std::vector &result, size_t num_elements, + int world_size, int rank) { + bool correct = true; + for (int src_rank = 0; src_rank < world_size; ++src_rank) { + const size_t offset = static_cast(src_rank) * num_elements; + const bool block_correct = + Validator::ValidateResult(result.data() + offset, num_elements, + static_cast(src_rank + 1), rank); + correct = block_correct && correct; + } + return correct; +} + +void WorkerThread(ThreadArgs args) { + constexpr Device::Type kDevType = + ListGetBest(EnabledDevices{}); + using Rt = Runtime; + + CHECK_RT(Rt, Rt::SetDevice(args.rank)); + + std::array hostname{}; + if (gethostname(hostname.data(), hostname.size()) != 0) { + std::cerr << "Failed to query the hostname for the AllGather worker." + << std::endl; + std::exit(EXIT_FAILURE); + } + hostname.back() = '\0'; + + std::cout << "[Rank " << args.rank << "] Host: " << hostname.data() + << " | GPU: " << Device::StringFromType(kDevType) << " | Device " + << args.rank << std::endl; + + infinicclComm_t comm = nullptr; + CHECK_INFINI(infinicclCommInitRank(&comm, args.size, args.id, args.rank)); + + const size_t send_bytes = args.num_elements * sizeof(float); + const size_t recv_elements = + args.num_elements * static_cast(args.size); + const size_t recv_bytes = recv_elements * sizeof(float); + + std::vector h_send(args.num_elements, + static_cast(args.rank + 1)); + std::vector h_recv(recv_elements, 0.0f); + + float *d_send = nullptr; + float *d_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), send_bytes)); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_recv), recv_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), send_bytes, + Rt::MemcpyHostToDevice)); + + auto run_scenario = [&](bool in_place) { + std::fill(h_recv.begin(), h_recv.end(), 0.0f); + const size_t local_offset = + static_cast(args.rank) * args.num_elements; + if (in_place) { + std::fill_n(h_recv.begin() + local_offset, args.num_elements, + static_cast(args.rank + 1)); + } + + CHECK_RT(Rt, Rt::Memcpy(d_recv, h_recv.data(), recv_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + const void *active_send = + in_place ? static_cast(d_recv + local_offset) + : static_cast(d_send); + + for (int i = 0; i < args.warmup_iter; ++i) { + CHECK_INFINI(infinicclAllGather(active_send, d_recv, args.num_elements, + infinicclFloat32, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + Timer timer; + for (int i = 0; i < args.profile_iter; ++i) { + CHECK_INFINI(infinicclAllGather(active_send, d_recv, args.num_elements, + infinicclFloat32, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const double elapsed_ms = + timer.ElapsedMs() / static_cast(args.profile_iter); + + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, recv_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + const bool correct = + ValidateAllGather(h_recv, args.num_elements, args.size, args.rank); + ScenarioState *state = in_place ? args.in_place : args.out_of_place; + if (!correct) { + state->correct.store(false, std::memory_order_relaxed); + } + WaitForScenario(state, args.size); + + if (args.rank == 0) { + PrintResult(in_place ? "In-Place" : "Out-of-Place", + state->correct.load(std::memory_order_acquire), h_recv, + args.num_elements, args.size, elapsed_ms); + } + }; + + run_scenario(false); + run_scenario(true); + + CHECK_RT(Rt, Rt::Free(d_send)); + CHECK_RT(Rt, Rt::Free(d_recv)); + CHECK_INFINI(infinicclCommDestroy(comm)); +} + +void PrintUsage(const char *program) { + std::cout << "Usage: " << program << " [options]\n" + << "Options:\n" + << " -g Number of GPUs (default: 8)\n" + << " -w Warmup iterations (default: 2)\n" + << " -p Profile iterations (default: 20)\n" + << " -n Elements per rank (default: 1048576)\n"; +} + +} // namespace + +int main(int argc, char **argv) { + int num_gpus = 8; + int warmup_iters = 2; + int profile_iters = 20; + size_t num_elements = 1 << 20; + + int opt = 0; + while ((opt = getopt(argc, argv, "g:w:p:n:h")) != -1) { + bool parsed = false; + switch (opt) { + case 'g': + parsed = ParsePositiveNumber(optarg, &num_gpus); + break; + case 'w': + parsed = ParsePositiveNumber(optarg, &warmup_iters); + break; + case 'p': + parsed = ParsePositiveNumber(optarg, &profile_iters); + break; + case 'n': + parsed = ParsePositiveNumber(optarg, &num_elements); + break; + case 'h': + PrintUsage(argv[0]); + return EXIT_SUCCESS; + default: + PrintUsage(argv[0]); + return EXIT_FAILURE; + } + + if (!parsed) { + std::cerr << "Invalid positive numeric option for AllGather." + << std::endl; + return EXIT_FAILURE; + } + } + + if (optind != argc) { + std::cerr << "Unexpected positional argument for AllGather." << std::endl; + return EXIT_FAILURE; + } + if (num_elements > std::numeric_limits::max() / sizeof(float) || + static_cast(num_gpus) > + std::numeric_limits::max() / num_elements || + num_elements * static_cast(num_gpus) > + std::numeric_limits::max() / sizeof(float)) { + std::cerr << "AllGather buffer size overflows `size_t`." << std::endl; + return EXIT_FAILURE; + } + + std::array hostname{}; + if (gethostname(hostname.data(), hostname.size()) != 0) { + std::cerr << "Failed to query the hostname for AllGather." << std::endl; + return EXIT_FAILURE; + } + hostname.back() = '\0'; + std::cout << "[Main Process] Host: " << hostname.data() + << " | Target GPUs: " << num_gpus << std::endl; + + infinicclUniqueId shared_id{}; + CHECK_INFINI(infinicclGetUniqueId(&shared_id)); + + ScenarioState out_of_place; + ScenarioState in_place; + std::vector 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, &out_of_place, &in_place}; + threads.emplace_back(WorkerThread, args); + } + + for (auto &thread : threads) { + if (thread.joinable()) { + thread.join(); + } + } + + const bool correct = out_of_place.correct.load(std::memory_order_acquire) && + in_place.correct.load(std::memory_order_acquire); + if (correct) { + std::cout << "[Main Process] All AllGather scenarios passed." << std::endl; + } else { + std::cerr << "[Main Process] AllGather validation failed." << std::endl; + } + std::cout + << "[Main Process] All worker threads joined. InfiniCCL finalized safely." + << std::endl; + return correct ? EXIT_SUCCESS : EXIT_FAILURE; +} diff --git a/examples/ccl_mpi_hybrid/all_gather.cc b/examples/ccl_mpi_hybrid/all_gather.cc new file mode 100644 index 0000000..b44317d --- /dev/null +++ b/examples/ccl_mpi_hybrid/all_gather.cc @@ -0,0 +1,293 @@ +/** + * InfiniCCL Example: AllGather (OpenMPI + CCL Hybrid) + * + * This example first uses AllGather through an OpenMPI inter communicator to + * distribute a native CCL unique ID. It then initializes a native CCL + * communicator and validates out-of-place and in-place GPU AllGather. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// 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; + +namespace { + +bool ParseLocalRank(const char *text, int *local_rank) { + if (!text || !local_rank) { + return false; + } + + int parsed = -1; + const char *end = text + std::strlen(text); + const auto result = std::from_chars(text, end, parsed); + if (result.ec != std::errc{} || result.ptr != end || parsed < 0) { + return false; + } + + *local_rank = parsed; + return true; +} + +bool ValidateAllGather(const std::vector &result, size_t num_elements, + int world_size, int rank) { + bool correct = true; + for (int src_rank = 0; src_rank < world_size; ++src_rank) { + const size_t offset = static_cast(src_rank) * num_elements; + const bool block_correct = + Validator::ValidateResult(result.data() + offset, num_elements, + static_cast(src_rank + 1), rank); + correct = block_correct && correct; + } + return correct; +} + +void PrintResult(const char *scenario, bool correct, + const std::vector &result, size_t num_elements, + int world_size, double elapsed_ms) { + constexpr const char *kGreen = "\033[32m"; + constexpr const char *kRed = "\033[31m"; + constexpr const char *kReset = "\033[0m"; + + std::cout << "\n=== " << scenario + << " Hybrid CCL AllGather Results ===" << std::endl; + std::cout << "Correct: " + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) + << std::endl; + std::cout << "Average time: " << elapsed_ms << " ms" << std::endl; + std::cout << "Sample blocks: "; + for (int src_rank = 0; src_rank < world_size && src_rank < 4; ++src_rank) { + const size_t offset = static_cast(src_rank) * num_elements; + std::cout << "[r" << src_rank << ": " << result[offset] << "] "; + } + std::cout << std::endl; +} + +bool RunAllGatherExample(int argc, char **argv) { + constexpr Device::Type kDevType = + ListGetBest(EnabledDevices{}); + using Rt = Runtime; + + constexpr int kWarmupIterations = 2; + constexpr int kProfileIterations = 20; + constexpr size_t kNumElements = 1 << 20; + + CHECK_INFINI(infinicclInit(&argc, &argv)); + + int rank = -1; + int size = 0; + CHECK_INFINI(infinicclGetRank(&rank)); + CHECK_INFINI(infinicclGetSize(&size)); + if (size <= 0) { + std::cerr << "Invalid world size for hybrid AllGather." << std::endl; + std::exit(EXIT_FAILURE); + } + + int local_rank = -1; + if (!ParseLocalRank(std::getenv("OMPI_COMM_WORLD_LOCAL_RANK"), &local_rank)) { + std::cerr << "Missing or invalid `OMPI_COMM_WORLD_LOCAL_RANK`." + << std::endl; + std::exit(EXIT_FAILURE); + } + CHECK_RT(Rt, Rt::SetDevice(local_rank)); + + std::array hostname{}; + if (gethostname(hostname.data(), hostname.size()) != 0) { + std::cerr << "Failed to query the hostname for hybrid AllGather." + << std::endl; + std::exit(EXIT_FAILURE); + } + hostname.back() = '\0'; + std::cout << "[Rank " << rank << "] Host: " << hostname.data() + << " | GPU: " << Device::StringFromType(kDevType) << " | Device " + << local_rank << std::endl; + + infinicclComm_t comm = nullptr; + CHECK_INFINI(infinicclCommInitAll(&comm, size, nullptr)); + + // Bootstrap the native communicator with AllGather itself. At this point + // only the OpenMPI inter communicator exists, so the CCL provider delegates + // this call to the existing OpenMPI staging implementation. + infinicclUniqueId id{}; + if (rank == 0) { + CHECK_INFINI(infinicclGetUniqueId(&id)); + } + const size_t id_bytes = sizeof(id); + if (static_cast(size) > + std::numeric_limits::max() / id_bytes) { + std::cerr << "AllGather bootstrap buffer size overflows `size_t`." + << std::endl; + std::exit(EXIT_FAILURE); + } + const size_t gathered_id_bytes = id_bytes * static_cast(size); + + unsigned char *d_id_send = nullptr; + unsigned char *d_id_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_id_send), id_bytes)); + CHECK_RT( + Rt, Rt::Malloc(reinterpret_cast(&d_id_recv), gathered_id_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_id_send, &id, id_bytes, Rt::MemcpyHostToDevice)); + CHECK_INFINI(infinicclAllGather(d_id_send, d_id_recv, id_bytes, infinicclChar, + comm, nullptr)); + CHECK_RT(Rt, Rt::Memcpy(&id, d_id_recv, id_bytes, Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + CHECK_RT(Rt, Rt::Free(d_id_send)); + CHECK_RT(Rt, Rt::Free(d_id_recv)); + + CHECK_INFINI(infinicclCommInitRank(&comm, size, id, rank)); + + if (kNumElements > std::numeric_limits::max() / sizeof(float) || + static_cast(size) > + std::numeric_limits::max() / kNumElements || + kNumElements * static_cast(size) > + std::numeric_limits::max() / sizeof(float)) { + std::cerr << "Hybrid AllGather buffer size overflows `size_t`." + << std::endl; + std::exit(EXIT_FAILURE); + } + + const size_t send_bytes = kNumElements * sizeof(float); + const size_t recv_elements = kNumElements * static_cast(size); + const size_t recv_bytes = recv_elements * sizeof(float); + std::vector h_send(kNumElements, static_cast(rank + 1)); + std::vector h_recv(recv_elements, 0.0f); + + float *d_send = nullptr; + float *d_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), send_bytes)); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_recv), recv_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), send_bytes, + Rt::MemcpyHostToDevice)); + + std::array local_correct{true, true}; + std::array elapsed_ms{}; + for (int scenario = 0; scenario < 2; ++scenario) { + const bool in_place = scenario == 1; + std::fill(h_recv.begin(), h_recv.end(), 0.0f); + const size_t local_offset = static_cast(rank) * kNumElements; + if (in_place) { + std::fill_n(h_recv.begin() + local_offset, kNumElements, + static_cast(rank + 1)); + } + + CHECK_RT(Rt, Rt::Memcpy(d_recv, h_recv.data(), recv_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const void *active_send = + in_place ? static_cast(d_recv + local_offset) + : static_cast(d_send); + + for (int i = 0; i < kWarmupIterations; ++i) { + CHECK_INFINI(infinicclAllGather(active_send, d_recv, kNumElements, + infinicclFloat32, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + Timer timer; + for (int i = 0; i < kProfileIterations; ++i) { + CHECK_INFINI(infinicclAllGather(active_send, d_recv, kNumElements, + infinicclFloat32, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + elapsed_ms[scenario] = + timer.ElapsedMs() / static_cast(kProfileIterations); + + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, recv_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + local_correct[scenario] = + ValidateAllGather(h_recv, kNumElements, size, rank); + } + + // Gather both local validation flags so rank 0 reports cluster-wide results + // and every rank reaches communicator destruction only after validation. + std::array h_status_send{local_correct[0] ? 1 : 0, + local_correct[1] ? 1 : 0}; + std::vector h_status_recv(static_cast(size) * 2, 0); + int32_t *d_status_send = nullptr; + int32_t *d_status_recv = nullptr; + const size_t status_send_bytes = + h_status_send.size() * sizeof(h_status_send[0]); + const size_t status_recv_bytes = + h_status_recv.size() * sizeof(h_status_recv[0]); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_status_send), + status_send_bytes)); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_status_recv), + status_recv_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_status_send, h_status_send.data(), + status_send_bytes, Rt::MemcpyHostToDevice)); + CHECK_INFINI(infinicclAllGather(d_status_send, d_status_recv, + h_status_send.size(), infinicclInt32, comm, + nullptr)); + CHECK_RT(Rt, Rt::Memcpy(h_status_recv.data(), d_status_recv, + status_recv_bytes, Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + std::array global_correct{true, true}; + for (int src_rank = 0; src_rank < size; ++src_rank) { + for (int scenario = 0; scenario < 2; ++scenario) { + const size_t index = static_cast(src_rank) * 2 + scenario; + global_correct[scenario] = + global_correct[scenario] && h_status_recv[index] == 1; + } + } + + if (rank == 0) { + // Rank 0's gathered payload is representative after cluster-wide status + // aggregation because every rank validates the complete result. + PrintResult("Out-of-Place", global_correct[0], h_recv, kNumElements, size, + elapsed_ms[0]); + PrintResult("In-Place", global_correct[1], h_recv, kNumElements, size, + elapsed_ms[1]); + } + + CHECK_RT(Rt, Rt::Free(d_status_send)); + CHECK_RT(Rt, Rt::Free(d_status_recv)); + CHECK_RT(Rt, Rt::Free(d_send)); + CHECK_RT(Rt, Rt::Free(d_recv)); + CHECK_INFINI(infinicclCommDestroy(comm)); + CHECK_INFINI(infinicclFinalize()); + + if (rank == 0) { + if (global_correct[0] && global_correct[1]) { + std::cout << "[Main Process] All hybrid AllGather scenarios passed." + << std::endl; + } else { + std::cerr << "[Main Process] Hybrid AllGather validation failed." + << std::endl; + } + std::cout << "InfiniCCL finalized." << std::endl; + } + return global_correct[0] && global_correct[1]; +} + +} // namespace + +int main(int argc, char **argv) { + return RunAllGatherExample(argc, argv) ? EXIT_SUCCESS : EXIT_FAILURE; +} diff --git a/examples/mpi/all_gather.cc b/examples/mpi/all_gather.cc index c0f1a62..a503280 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); + const bool block_correct = Validator::ValidateResult( + h_recv.data() + offset, kNumElements, expected, rank); + correct = block_correct && 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,8 @@ 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 +167,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..c81e39b --- /dev/null +++ b/src/backends/ccl/common/impl/all_gather.h @@ -0,0 +1,63 @@ +#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 +struct DeferredAllGather { + using type = AllGather; +}; + +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; + + const bool has_native_comm = comm && comm->intra_comm() && + comm->intra_comm_backend() == backend && + comm->device_type() == device; + if (!has_native_comm) { + if (!comm || !comm->inter_comm() || + comm->inter_comm_backend() != BackendType::kOmpi) { + return ReturnStatus::kInternalError; + } + + using FallbackOperation = typename DeferredAllGather::type; + if constexpr (BackendEnabled::value) { + return AllGatherImpl::Apply( + send_buff, recv_buff, count, data_type, comm, stream); + } + + return ReturnStatus::kInternalError; + } + + auto *instance = static_cast(comm->intra_comm()); + if (!instance->handle) { + return ReturnStatus::kInternalError; + } + + typename Api::DataType native_type{}; + if (!TypeMap::ToBackendDataType(data_type, &native_type)) { + return ReturnStatus::kNotSupported; + } + + return Api::Check(Api::AllGather( + send_buff, recv_buff, count, native_type, instance->handle, + reinterpret_cast(stream))); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_GATHER_H_ diff --git a/src/backends/ccl/mccl/api.h b/src/backends/ccl/mccl/api.h index a5d2bcd..0f70596 100644 --- a/src/backends/ccl/mccl/api.h +++ b/src/backends/ccl/mccl/api.h @@ -49,6 +49,13 @@ struct McclApi { return mcclAllReduce(send_buff, recv_buff, count, data_type, op, comm, stream); } + + static Result AllGather(const void *send_buff, void *recv_buff, + size_t send_count, DataType data_type, Comm comm, + Stream stream) { + return mcclAllGather(send_buff, recv_buff, send_count, data_type, comm, + stream); + } }; } // namespace infini::ccl diff --git a/src/backends/ccl/mccl/impl/all_gather.h b/src/backends/ccl/mccl/impl/all_gather.h new file mode 100644 index 0000000..65dc9b8 --- /dev/null +++ b/src/backends/ccl/mccl/impl/all_gather.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_ALL_GATHER_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_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_MCCL_IMPL_ALL_GATHER_H_ diff --git a/src/backends/ccl/nccl/api.h b/src/backends/ccl/nccl/api.h index e7b6119..6ece5ad 100644 --- a/src/backends/ccl/nccl/api.h +++ b/src/backends/ccl/nccl/api.h @@ -46,6 +46,13 @@ struct NcclApi { return ncclAllReduce(send_buff, recv_buff, count, data_type, op, comm, 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); + } }; } // namespace infini::ccl 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/backends/mpi/ompi/impl/all_gather.h b/src/backends/mpi/ompi/impl/all_gather.h index a1aee69..a8e75b9 100644 --- a/src/backends/mpi/ompi/impl/all_gather.h +++ b/src/backends/mpi/ompi/impl/all_gather.h @@ -1,11 +1,14 @@ #ifndef INFINI_CCL_BACKENDS_MPI_OMPI_IMPL_ALL_GATHER_H_ #define INFINI_CCL_BACKENDS_MPI_OMPI_IMPL_ALL_GATHER_H_ +#include +#include + #include "backends/mpi/ompi/checks.h" #include "backends/mpi/ompi/comm_instance.h" -#include "backends/mpi/ompi/type_map.h" #include "base/all_gather.h" #include "communicator.h" +#include "data_type_impl.h" #include "dispatcher.h" #include "logging.h" @@ -28,19 +31,38 @@ class AllGatherImpl { return ReturnStatus::kInternalError; } - MPI_Datatype mpi_type = DataTypeToOmpiType(data_type); + if (comm->size() <= 0) { + LOG("Invalid world size for `AllGather`."); + return ReturnStatus::kInternalError; + } + size_t type_size = kDataTypeToSize.at(data_type); + if (count > std::numeric_limits::max() / type_size) { + LOG("Byte size overflow for `AllGather`."); + return ReturnStatus::kInvalidArgument; + } size_t send_bytes = count * type_size; - size_t recv_count = count * static_cast(comm->size()); - size_t recv_bytes = recv_count * type_size; + if (send_bytes > static_cast(std::numeric_limits::max())) { + LOG("Per-rank byte count exceeds MPI int range for `AllGather`."); + return ReturnStatus::kInvalidArgument; + } + + size_t world_size = static_cast(comm->size()); + if (world_size != 0 && + send_bytes > std::numeric_limits::max() / world_size) { + LOG("Receive byte size overflow for `AllGather`."); + return ReturnStatus::kInvalidArgument; + } + size_t recv_bytes = send_bytes * world_size; + int mpi_byte_count = static_cast(send_bytes); // Handle GPU Memory (Staging Pattern) // Note: we simply use host-staging for now. - void *host_sendbuf = malloc(send_bytes); - void *host_recvbuf = malloc(recv_bytes); + void *host_sendbuf = std::malloc(send_bytes == 0 ? 1 : send_bytes); + void *host_recvbuf = std::malloc(recv_bytes == 0 ? 1 : recv_bytes); if (!host_sendbuf || !host_recvbuf) { - free(host_sendbuf); - free(host_recvbuf); + std::free(host_sendbuf); + std::free(host_recvbuf); LOG("Failed to allocate host buffers for `AllGather` staging."); return ReturnStatus::kSystemError; } @@ -50,14 +72,15 @@ class AllGatherImpl { CHECK_STATUS(Rt, Rt::StreamSynchronize(static_cast(stream))); - INFINI_CHECK_MPI(MPI_Allgather(host_sendbuf, count, mpi_type, host_recvbuf, - count, mpi_type, inst->handle)); + INFINI_CHECK_MPI(MPI_Allgather(host_sendbuf, mpi_byte_count, MPI_BYTE, + host_recvbuf, mpi_byte_count, MPI_BYTE, + inst->handle)); CHECK_STATUS(Rt, Rt::Memcpy(recv_buff, host_recvbuf, recv_bytes, Rt::MemcpyHostToDevice)); - free(host_sendbuf); - free(host_recvbuf); + std::free(host_sendbuf); + std::free(host_recvbuf); return ReturnStatus::kSuccess; } diff --git a/src/base/all_gather.h b/src/base/all_gather.h index c0072d8..735045e 100644 --- a/src/base/all_gather.h +++ b/src/base/all_gather.h @@ -19,9 +19,13 @@ 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 (HasInvalidArgs(send_buff, recv_buff, count, datatype, comm_handle)) { return ReturnStatus::kInvalidArgument; } + if (count == 0) { + return ReturnStatus::kSuccess; + } + auto *comm = static_cast(comm_handle); return AllGatherImpl::Apply( send_buff, recv_buff, count, datatype, comm, stream); @@ -29,20 +33,24 @@ class AllGather : public Operation { private: static bool HasInvalidArgs(const void *send_buff, void *recv_buff, - DataType datatype, void *comm_handle) { + size_t count, DataType datatype, + void *comm_handle) { if (!comm_handle) { // TODO(lzm): change to use `glog`. LOG("Invalid communicator handle for `AllGather`."); return true; } - 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; } + if (count == 0) { + return false; + } + if (!send_buff || !recv_buff) { + LOG("Invalid buffer pointer for `AllGather`."); + return true; + } return false; } };