From 773a7d32846a547c83effb5f0db9f5d8a208424e Mon Sep 17 00:00:00 2001 From: GordonYang1 <1468121796@qq.com> Date: Sat, 22 Aug 2026 10:39:57 +0800 Subject: [PATCH] feat: support CCL `Gather` --- examples/ccl/gather.cc | 331 ++++++++++++++++++++++++++ examples/ccl_mpi_hybrid/gather.cc | 312 ++++++++++++++++++++++++ examples/mpi/gather.cc | 273 +++++++++++++-------- src/backends/ccl/common/impl/gather.h | 108 +++++++++ src/backends/ccl/mccl/api.h | 14 ++ src/backends/ccl/mccl/impl/gather.h | 17 ++ src/backends/ccl/nccl/api.h | 14 ++ src/backends/ccl/nccl/impl/gather.h | 17 ++ src/backends/mpi/ompi/impl/gather.h | 51 ++-- src/base/gather.h | 32 +-- 10 files changed, 1036 insertions(+), 133 deletions(-) create mode 100644 examples/ccl/gather.cc create mode 100644 examples/ccl_mpi_hybrid/gather.cc create mode 100644 src/backends/ccl/common/impl/gather.h create mode 100644 src/backends/ccl/mccl/impl/gather.h create mode 100644 src/backends/ccl/nccl/impl/gather.h diff --git a/examples/ccl/gather.cc b/examples/ccl/gather.cc new file mode 100644 index 0000000..9cb549a --- /dev/null +++ b/examples/ccl/gather.cc @@ -0,0 +1,331 @@ +/** + * InfiniCCL Example: Thread-per-GPU Single-Node Gather + * + * This example creates one native CCL rank per GPU and gathers one block from + * every rank into rank 0 using grouped point-to-point operations. + */ + +#include + +#include +#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 { + +constexpr int kRoot = 0; + +struct ScenarioState { + std::atomic correct{true}; + std::atomic completed{0}; +}; + +struct ThreadArgs { + int rank; + int size; + infinicclUniqueId id; + size_t num_elements; + int warmup_iterations; + int profile_iterations; + ScenarioState *state; +}; + +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; +} + +bool ValidateGather(const std::vector &result, size_t num_elements, + int world_size) { + bool correct = true; + for (int source = 0; source < world_size; ++source) { + const size_t offset = static_cast(source) * num_elements; + const bool block_correct = + Validator::ValidateResult(result.data() + offset, num_elements, + static_cast(source + 1), kRoot); + correct = block_correct && correct; + } + return correct; +} + +void PrintGatherMetrics(size_t num_elements, int world_size, + double elapsed_ms) { + constexpr double kBytesPerMiB = 1024.0 * 1024.0; + constexpr double kBytesPerGB = 1.0e9; + const double rank_bytes = static_cast(num_elements) * sizeof(float); + const double gathered_bytes = rank_bytes * static_cast(world_size); + const auto original_flags = std::cout.flags(); + const auto original_precision = std::cout.precision(); + + std::cout << "Data size per rank: " << num_elements << " floats (" + << std::fixed << std::setprecision(2) << rank_bytes / kBytesPerMiB + << " MiB)" << std::endl; + std::cout << "Total data at root: " + << num_elements * static_cast(world_size) << " floats (" + << gathered_bytes / kBytesPerMiB << " MiB)" << std::endl; + std::cout << "Time: " << std::setprecision(3) << elapsed_ms << " ms" + << std::endl; + if (elapsed_ms > 0.0 && std::isfinite(elapsed_ms)) { + const double algorithm_bandwidth = + gathered_bytes / kBytesPerGB / (elapsed_ms / 1000.0); + const double bus_bandwidth = algorithm_bandwidth * + static_cast(world_size - 1) / + static_cast(world_size); + std::cout << "Throughput: " << std::setprecision(2) << bus_bandwidth + << " GB/s (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: " << algorithm_bandwidth << " GB/s" + << std::endl; + } else { + std::cout << "Throughput: N/A (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: N/A" << std::endl; + } + + std::cout.flags(original_flags); + std::cout.precision(original_precision); +} + +void WaitForAll(ScenarioState *state, int world_size) { + state->completed.fetch_add(1, std::memory_order_acq_rel); + while (state->completed.load(std::memory_order_acquire) < world_size) { + std::this_thread::yield(); + } +} + +void PrintResult(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=== CCL Gather Results ===" << std::endl; + std::cout << "Correct: " + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) + << std::endl; + std::cout << "Root rank: " << kRoot << std::endl; + std::cout << "Sample receive blocks: "; + for (int source = 0; source < std::min(world_size, 4); ++source) { + const size_t offset = static_cast(source) * num_elements; + std::cout << "[r" << source << ": " << result[offset] << "] "; + } + std::cout << std::endl; + PrintGatherMetrics(num_elements, world_size, elapsed_ms); +} + +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 Gather 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 rank_bytes = args.num_elements * sizeof(float); + const size_t gathered_elements = + args.num_elements * static_cast(args.size); + const size_t gathered_bytes = gathered_elements * sizeof(float); + std::vector h_send(args.num_elements, + static_cast(args.rank + 1)); + std::vector h_recv; + if (args.rank == kRoot) { + h_recv.resize(gathered_elements, 0.0f); + } + + float *d_send = nullptr; + float *d_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), rank_bytes)); + if (args.rank == kRoot) { + CHECK_RT(Rt, + Rt::Malloc(reinterpret_cast(&d_recv), gathered_bytes)); + } + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), rank_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + for (int i = 0; i < args.warmup_iterations; ++i) { + CHECK_INFINI(infinicclGather(d_send, d_recv, args.num_elements, + infinicclFloat32, kRoot, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + Timer timer; + for (int i = 0; i < args.profile_iterations; ++i) { + CHECK_INFINI(infinicclGather(d_send, d_recv, args.num_elements, + infinicclFloat32, kRoot, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const double elapsed_ms = + timer.ElapsedMs() / static_cast(args.profile_iterations); + + if (args.rank == kRoot) { + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, gathered_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + args.state->correct.store( + ValidateGather(h_recv, args.num_elements, args.size), + std::memory_order_release); + } + + WaitForAll(args.state, args.size); + if (args.rank == kRoot) { + PrintResult(args.state->correct.load(std::memory_order_acquire), h_recv, + args.num_elements, args.size, elapsed_ms); + } + + CHECK_RT(Rt, Rt::Free(d_send)); + if (args.rank == kRoot) { + 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 contributed by each rank " + "(default: 1048576)\n"; +} + +} // namespace + +int main(int argc, char **argv) { + int num_gpus = 8; + int warmup_iterations = 2; + int profile_iterations = 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_iterations); + break; + case 'p': + parsed = ParsePositiveNumber(optarg, &profile_iterations); + 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 Gather." << std::endl; + return EXIT_FAILURE; + } + } + + if (optind != argc) { + std::cerr << "Unexpected positional argument for Gather." << std::endl; + return EXIT_FAILURE; + } + if (static_cast(num_gpus) > + std::numeric_limits::max() / num_elements || + num_elements * static_cast(num_gpus) > + std::numeric_limits::max() / sizeof(float)) { + std::cerr << "Gather 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 Gather." << std::endl; + return EXIT_FAILURE; + } + hostname.back() = '\0'; + std::cout << "[Main Process] Host: " << hostname.data() + << " | Target GPUs: " << num_gpus << std::endl; + std::cout << "[Main Process] Elements per rank: " << num_elements + << " floats | Warmup: " << warmup_iterations + << " | Profile: " << profile_iterations << std::endl; + + infinicclUniqueId shared_id{}; + CHECK_INFINI(infinicclGetUniqueId(&shared_id)); + + ScenarioState state; + 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_iterations, profile_iterations, + &state}; + threads.emplace_back(WorkerThread, args); + } + + for (auto &thread : threads) { + if (thread.joinable()) { + thread.join(); + } + } + + const bool correct = state.correct.load(std::memory_order_acquire); + if (correct) { + std::cout << "[Main Process] CCL Gather validation passed." << std::endl; + } else { + std::cerr << "[Main Process] CCL Gather 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/gather.cc b/examples/ccl_mpi_hybrid/gather.cc new file mode 100644 index 0000000..131ee2a --- /dev/null +++ b/examples/ccl_mpi_hybrid/gather.cc @@ -0,0 +1,312 @@ +/** + * InfiniCCL Example: Gather (OpenMPI + CCL Hybrid) + * + * This example first exercises Gather through its OpenMPI fallback, then + * initializes a native CCL communicator and profiles the grouped-P2P path. + */ + +#include + +#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" +#include "device.h" +#include "runtime.h" +#include "traits.h" + +using namespace infini::ccl; + +namespace { + +constexpr int kRoot = 0; + +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 ValidateGather(const std::vector &result, size_t num_elements, + int world_size) { + bool correct = true; + for (int source = 0; source < world_size; ++source) { + const size_t offset = static_cast(source) * num_elements; + const bool block_correct = + Validator::ValidateResult(result.data() + offset, num_elements, + static_cast(source + 1), kRoot); + correct = block_correct && correct; + } + return correct; +} + +void PrintGatherMetrics(size_t num_elements, int world_size, + double elapsed_ms) { + constexpr double kBytesPerMiB = 1024.0 * 1024.0; + constexpr double kBytesPerGB = 1.0e9; + const double rank_bytes = static_cast(num_elements) * sizeof(float); + const double gathered_bytes = rank_bytes * static_cast(world_size); + const auto original_flags = std::cout.flags(); + const auto original_precision = std::cout.precision(); + + std::cout << "Data size per rank: " << num_elements << " floats (" + << std::fixed << std::setprecision(2) << rank_bytes / kBytesPerMiB + << " MiB)" << std::endl; + std::cout << "Total data at root: " + << num_elements * static_cast(world_size) << " floats (" + << gathered_bytes / kBytesPerMiB << " MiB)" << std::endl; + std::cout << "Time: " << std::setprecision(3) << elapsed_ms << " ms" + << std::endl; + if (elapsed_ms > 0.0 && std::isfinite(elapsed_ms)) { + const double algorithm_bandwidth = + gathered_bytes / kBytesPerGB / (elapsed_ms / 1000.0); + const double bus_bandwidth = algorithm_bandwidth * + static_cast(world_size - 1) / + static_cast(world_size); + std::cout << "Throughput: " << std::setprecision(2) << bus_bandwidth + << " GB/s (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: " << algorithm_bandwidth << " GB/s" + << std::endl; + } else { + std::cout << "Throughput: N/A (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: N/A" << std::endl; + } + + std::cout.flags(original_flags); + std::cout.precision(original_precision); +} + +void PrintResult(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=== Hybrid CCL Gather Results ===" << std::endl; + std::cout << "Correct: " + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) + << std::endl; + std::cout << "Root rank: " << kRoot << std::endl; + std::cout << "Sample receive blocks: "; + for (int source = 0; source < std::min(world_size, 4); ++source) { + const size_t offset = static_cast(source) * num_elements; + std::cout << "[r" << source << ": " << result[offset] << "] "; + } + std::cout << std::endl; + PrintGatherMetrics(num_elements, world_size, elapsed_ms); +} + +bool RunGatherExample(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 Gather." << 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 Gather." << 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)); + + // Before a native communicator exists, the selected CCL Gather provider + // delegates this rank token exchange to the OpenMPI inter communicator. + const float bootstrap_value = static_cast(rank + 1); + std::vector h_bootstrap_recv; + if (rank == kRoot) { + h_bootstrap_recv.resize(static_cast(size), 0.0f); + } + float *d_bootstrap_send = nullptr; + float *d_bootstrap_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_bootstrap_send), + sizeof(float))); + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_bootstrap_recv), + static_cast(size) * sizeof(float))); + } + CHECK_RT(Rt, Rt::Memcpy(d_bootstrap_send, &bootstrap_value, sizeof(float), + Rt::MemcpyHostToDevice)); + CHECK_INFINI(infinicclGather(d_bootstrap_send, d_bootstrap_recv, 1, + infinicclFloat32, kRoot, comm, nullptr)); + + int32_t bootstrap_status = 1; + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Memcpy(h_bootstrap_recv.data(), d_bootstrap_recv, + static_cast(size) * sizeof(float), + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + for (int source = 0; source < size; ++source) { + if (h_bootstrap_recv[static_cast(source)] != + static_cast(source + 1)) { + bootstrap_status = 0; + } + } + } + CHECK_INFINI(infinicclBroadcast(&bootstrap_status, &bootstrap_status, 1, + infinicclInt32, kRoot, comm, nullptr)); + CHECK_RT(Rt, Rt::Free(d_bootstrap_send)); + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Free(d_bootstrap_recv)); + } + + if (bootstrap_status != 1) { + if (rank == kRoot) { + std::cerr << "OpenMPI Gather fallback validation failed." << std::endl; + } + CHECK_INFINI(infinicclCommDestroy(comm)); + CHECK_INFINI(infinicclFinalize()); + return false; + } + if (rank == kRoot) { + std::cout << "OpenMPI Gather fallback validation passed." << std::endl; + } + + infinicclUniqueId id{}; + if (rank == kRoot) { + CHECK_INFINI(infinicclGetUniqueId(&id)); + } + CHECK_INFINI(infinicclBroadcast(&id, &id, sizeof(id), infinicclUInt8, kRoot, + comm, nullptr)); + CHECK_INFINI(infinicclCommInitRank(&comm, size, id, rank)); + + const size_t world_size = static_cast(size); + if (kNumElements > std::numeric_limits::max() / world_size || + kNumElements * world_size > + std::numeric_limits::max() / sizeof(float)) { + std::cerr << "Hybrid Gather buffer size overflows `size_t`." << std::endl; + std::exit(EXIT_FAILURE); + } + const size_t rank_bytes = kNumElements * sizeof(float); + const size_t gathered_elements = kNumElements * world_size; + const size_t gathered_bytes = gathered_elements * sizeof(float); + std::vector h_send(kNumElements, static_cast(rank + 1)); + std::vector h_recv; + if (rank == kRoot) { + h_recv.resize(gathered_elements, 0.0f); + } + + float *d_send = nullptr; + float *d_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), rank_bytes)); + if (rank == kRoot) { + CHECK_RT(Rt, + Rt::Malloc(reinterpret_cast(&d_recv), gathered_bytes)); + } + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), rank_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + for (int i = 0; i < kWarmupIterations; ++i) { + CHECK_INFINI(infinicclGather(d_send, d_recv, kNumElements, infinicclFloat32, + kRoot, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + Timer timer; + for (int i = 0; i < kProfileIterations; ++i) { + CHECK_INFINI(infinicclGather(d_send, d_recv, kNumElements, infinicclFloat32, + kRoot, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const double elapsed_ms = + timer.ElapsedMs() / static_cast(kProfileIterations); + + int32_t validation_status = 1; + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, gathered_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + validation_status = + ValidateGather(h_recv, kNumElements, size) ? int32_t{1} : int32_t{0}; + } + CHECK_INFINI(infinicclBroadcast(&validation_status, &validation_status, 1, + infinicclInt32, kRoot, comm, nullptr)); + const bool correct = validation_status == 1; + + if (rank == kRoot) { + PrintResult(correct, h_recv, kNumElements, size, elapsed_ms); + } + + CHECK_RT(Rt, Rt::Free(d_send)); + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Free(d_recv)); + } + CHECK_INFINI(infinicclCommDestroy(comm)); + CHECK_INFINI(infinicclFinalize()); + + if (rank == kRoot) { + if (correct) { + std::cout << "[Main Process] Hybrid CCL Gather validation passed." + << std::endl; + } else { + std::cerr << "[Main Process] Hybrid CCL Gather validation failed." + << std::endl; + } + std::cout << "InfiniCCL finalized." << std::endl; + } + return correct; +} + +} // namespace + +int main(int argc, char **argv) { + return RunGatherExample(argc, argv) ? EXIT_SUCCESS : EXIT_FAILURE; +} diff --git a/examples/mpi/gather.cc b/examples/mpi/gather.cc index b488df9..4858232 100644 --- a/examples/mpi/gather.cc +++ b/examples/mpi/gather.cc @@ -1,13 +1,24 @@ /** - * InfiniCCL Example: Gather - * * This example demonstrates the API for performing a collective - * data gathering across multiple GPUs and nodes, where only `root` - * receives the gathered result. + * InfiniCCL Example: Gather (MPI Backend) + * + * Every rank contributes one GPU-resident block. Rank 0 receives the blocks + * in rank order, validates them, and reports Gather-specific bandwidth. */ #include +#include +#include +#include +#include +#include +#include +#include +#include #include +#include +#include +#include #include // Public API @@ -25,146 +36,216 @@ using namespace infini::ccl; -void RunGatherExample(int argc, char **argv, int warmup_iter, int profile_iter, - const size_t kNumElements) { +namespace { + +constexpr int kRoot = 0; + +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 ValidateGather(const std::vector &result, size_t num_elements, + int world_size) { + bool correct = true; + for (int source = 0; source < world_size; ++source) { + const size_t offset = static_cast(source) * num_elements; + const bool block_correct = + Validator::ValidateResult(result.data() + offset, num_elements, + static_cast(source + 1), kRoot); + correct = block_correct && correct; + } + return correct; +} + +void PrintGatherMetrics(size_t num_elements, int world_size, + double elapsed_ms) { + constexpr double kBytesPerMiB = 1024.0 * 1024.0; + constexpr double kBytesPerGB = 1.0e9; + const double rank_bytes = static_cast(num_elements) * sizeof(float); + const double gathered_bytes = rank_bytes * static_cast(world_size); + const auto original_flags = std::cout.flags(); + const auto original_precision = std::cout.precision(); + + std::cout << "Data size per rank: " << num_elements << " floats (" + << std::fixed << std::setprecision(2) << rank_bytes / kBytesPerMiB + << " MiB)" << std::endl; + std::cout << "Total data at root: " + << num_elements * static_cast(world_size) << " floats (" + << gathered_bytes / kBytesPerMiB << " MiB)" << std::endl; + std::cout << "Time: " << std::setprecision(3) << elapsed_ms << " ms" + << std::endl; + if (elapsed_ms > 0.0 && std::isfinite(elapsed_ms)) { + const double algorithm_bandwidth = + gathered_bytes / kBytesPerGB / (elapsed_ms / 1000.0); + const double bus_bandwidth = algorithm_bandwidth * + static_cast(world_size - 1) / + static_cast(world_size); + std::cout << "Throughput: " << std::setprecision(2) << bus_bandwidth + << " GB/s (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: " << algorithm_bandwidth << " GB/s" + << std::endl; + } else { + std::cout << "Throughput: N/A (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: N/A" << std::endl; + } + + std::cout.flags(original_flags); + std::cout.precision(original_precision); +} + +bool RunGatherExample(int argc, char **argv, int warmup_iterations, + int profile_iterations, size_t num_elements) { constexpr Device::Type kDevType = ListGetBest(EnabledDevices{}); using Rt = Runtime; CHECK_INFINI(infinicclInit(&argc, &argv)); - int rank, size; + int rank = -1; + int size = 0; CHECK_INFINI(infinicclGetRank(&rank)); CHECK_INFINI(infinicclGetSize(&size)); + if (size <= 0) { + std::cerr << "Invalid world size for MPI Gather." << std::endl; + std::exit(EXIT_FAILURE); + } - 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); + 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::cout << "[Rank " << rank << "] Host: " << hostname - << " | GPU: " << Device::StringFromType(kDevType) << " " - << " | Device " << local_rank << std::endl; + std::array hostname{}; + if (gethostname(hostname.data(), hostname.size()) != 0) { + std::cerr << "Failed to query the hostname for MPI Gather." << 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; - // Setup Communicator infinicclComm_t comm = nullptr; CHECK_INFINI(infinicclCommInitAll(&comm, size, nullptr)); - // Root of the Gather - constexpr int kRoot = 0; - - // Prepare Data - std::vector h_send(kNumElements); - std::vector h_recv(kNumElements * size, 0.0f); - - // Initialize: each rank provides its (rank + 1) as data. - for (size_t i = 0; i < kNumElements; i++) { - h_send[i] = static_cast(rank + 1); + const size_t world_size = static_cast(size); + if (num_elements > std::numeric_limits::max() / world_size || + num_elements * world_size > + std::numeric_limits::max() / sizeof(float)) { + std::cerr << "MPI Gather buffer size overflows `size_t`." << std::endl; + std::exit(EXIT_FAILURE); } - - float *d_send, *d_recv; - size_t send_bytes = kNumElements * sizeof(*d_send); - size_t recv_bytes = send_bytes * 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)); - + const size_t rank_bytes = num_elements * sizeof(float); + const size_t gathered_elements = num_elements * world_size; + const size_t gathered_bytes = gathered_elements * sizeof(float); + std::vector h_send(num_elements, static_cast(rank + 1)); + std::vector h_recv; if (rank == kRoot) { - std::cout << "\n=== Performing Gather on GPU Memory ===" << std::endl; - std::cout << "Data size: " << kNumElements << " floats (" - << send_bytes / 1024 / 1024 << " MB)" << std::endl; - std::cout << "Operation: Gather" << std::endl; - std::cout << "Root Rank: " << kRoot << std::endl; - std::cout << "Warm-up iterations: " << warmup_iter << std::endl; - std::cout << "Profile iterations: " << profile_iter << std::endl; + h_recv.resize(gathered_elements, 0.0f); } + float *d_send = nullptr; + float *d_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), rank_bytes)); + if (rank == kRoot) { + CHECK_RT(Rt, + Rt::Malloc(reinterpret_cast(&d_recv), gathered_bytes)); + } + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), rank_bytes, + Rt::MemcpyHostToDevice)); CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); - // Warm-up - CHECK_INFINI(infinicclGather(d_send, d_recv, kNumElements, infinicclFloat32, - kRoot, comm, nullptr)); + if (rank == kRoot) { + std::cout << "\n=== Performing MPI Gather on GPU Memory ===" << std::endl; + std::cout << "Root rank: " << kRoot << std::endl; + std::cout << "Elements per rank: " << num_elements << " floats" + << std::endl; + std::cout << "Warm-up iterations: " << warmup_iterations << std::endl; + std::cout << "Profile iterations: " << profile_iterations << std::endl; + } - for (int i = 1; i < warmup_iter; ++i) { - CHECK_INFINI(infinicclGather(d_send, d_recv, kNumElements, infinicclFloat32, + for (int i = 0; i < warmup_iterations; ++i) { + CHECK_INFINI(infinicclGather(d_send, d_recv, num_elements, infinicclFloat32, kRoot, comm, nullptr)); } CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); - // Profiling Timer timer; - - for (int i = 0; i < profile_iter; i++) { - CHECK_INFINI(infinicclGather(d_send, d_recv, kNumElements, infinicclFloat32, + for (int i = 0; i < profile_iterations; ++i) { + CHECK_INFINI(infinicclGather(d_send, d_recv, num_elements, infinicclFloat32, kRoot, comm, nullptr)); } - CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); - double elapsed = timer.ElapsedMs() / static_cast(profile_iter); + const double elapsed_ms = + timer.ElapsedMs() / static_cast(profile_iterations); - // Result Validation (only meaningful on `root`). + int32_t validation_status = 1; if (rank == kRoot) { - CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, recv_bytes, + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, gathered_bytes, Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + validation_status = + ValidateGather(h_recv, num_elements, size) ? int32_t{1} : int32_t{0}; + } + CHECK_INFINI(infinicclBroadcast(&validation_status, &validation_status, 1, + infinicclInt32, kRoot, comm, nullptr)); + const bool correct = validation_status == 1; - bool correct = true; - - 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; - - correct &= Validator::ValidateResult(h_recv.data() + offset, kNumElements, - expected, rank); - } - - const char *GREEN = "\033[32m"; - const char *RED = "\033[31m"; - const char *RESET = "\033[0m"; - - std::cout << "\n=== Gather Results ===" << std::endl; + if (rank == kRoot) { + constexpr const char *kGreen = "\033[32m"; + constexpr const char *kRed = "\033[31m"; + constexpr const char *kReset = "\033[0m"; + std::cout << "\n=== MPI Gather Results ===" << std::endl; std::cout << "Correct: " - << (correct ? (GREEN + std::string("YES") + RESET) - : (RED + std::string("NO") + RESET)) + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) << std::endl; - - std::cout << "Sample blocks: "; - for (int src_rank = 0; src_rank < std::min(size, 4); ++src_rank) { - size_t offset = static_cast(src_rank) * kNumElements; - std::cout << "[r" << src_rank << ": " << h_recv[offset] << "] "; + std::cout << "Root rank: " << kRoot << std::endl; + std::cout << "Sample receive blocks: "; + for (int source = 0; source < std::min(size, 4); ++source) { + const size_t offset = static_cast(source) * num_elements; + std::cout << "[r" << source << ": " << h_recv[offset] << "] "; } std::cout << std::endl; - - Metrics metrics{elapsed, recv_bytes, size}; - metrics.Print(); + PrintGatherMetrics(num_elements, size, elapsed_ms); } - // Cleanup CHECK_RT(Rt, Rt::Free(d_send)); - CHECK_RT(Rt, Rt::Free(d_recv)); - + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Free(d_recv)); + } CHECK_INFINI(infinicclCommDestroy(comm)); CHECK_INFINI(infinicclFinalize()); if (rank == kRoot) { std::cout << "InfiniCCL finalized." << std::endl; } + return correct; } -int main(int argc, char **argv) { - int warmup_iters = 2; - int profile_iters = 20; - size_t num_elements = 1 << 20; - - RunGatherExample(argc, argv, warmup_iters, profile_iters, num_elements); +} // namespace - return EXIT_SUCCESS; +int main(int argc, char **argv) { + constexpr int kWarmupIterations = 2; + constexpr int kProfileIterations = 20; + constexpr size_t kNumElements = 1 << 20; + return RunGatherExample(argc, argv, kWarmupIterations, kProfileIterations, + kNumElements) + ? EXIT_SUCCESS + : EXIT_FAILURE; } diff --git a/src/backends/ccl/common/impl/gather.h b/src/backends/ccl/common/impl/gather.h new file mode 100644 index 0000000..5e2faa2 --- /dev/null +++ b/src/backends/ccl/common/impl/gather.h @@ -0,0 +1,108 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_GATHER_H_ +#define INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_GATHER_H_ + +#include +#include + +#include "backends/ccl/common/api.h" +#include "backends/ccl/common/comm_instance.h" +#include "base/gather.h" +#include "communicator.h" +#include "data_type_impl.h" +#include "logging.h" + +namespace infini::ccl { + +template +struct DeferredGather { + using type = Gather; +}; + +template +class CclGatherImpl { + public: + static ReturnStatus Apply(const void *send_buff, void *recv_buff, + size_t count, DataType data_type, int root, + 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 DeferredGather::type; + if constexpr (BackendEnabled::value) { + return GatherImpl::Apply( + send_buff, recv_buff, count, data_type, root, comm, stream); + } + + return ReturnStatus::kInternalError; + } + + if (comm->size() <= 0 || comm->rank() < 0 || comm->rank() >= comm->size() || + root < 0 || root >= comm->size()) { + LOG("Invalid rank, root, or world size for native CCL `Gather`."); + 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; + } + + const size_t type_size = kDataTypeToSize.at(data_type); + if (count > std::numeric_limits::max() / type_size) { + LOG("Per-rank byte size overflows `size_t` for native CCL `Gather`."); + return ReturnStatus::kInvalidArgument; + } + const size_t rank_bytes = count * type_size; + const size_t world_size = static_cast(comm->size()); + if (rank_bytes > std::numeric_limits::max() / world_size) { + LOG("Total byte size overflows `size_t` for native CCL `Gather`."); + return ReturnStatus::kInvalidArgument; + } + + auto native_stream = reinterpret_cast(stream); + ReturnStatus status = Api::Check(Api::GroupStart()); + if (status != ReturnStatus::kSuccess) { + return status; + } + + ReturnStatus first_error = Api::Check(Api::Send( + send_buff, count, native_type, root, instance->handle, native_stream)); + + if (comm->rank() == root) { + auto *recv_bytes = static_cast(recv_buff); + for (int peer = 0; peer < comm->size(); ++peer) { + const size_t offset = static_cast(peer) * rank_bytes; + status = Api::Check(Api::Recv(recv_bytes + offset, count, native_type, + peer, instance->handle, native_stream)); + if (first_error == ReturnStatus::kSuccess && + status != ReturnStatus::kSuccess) { + first_error = status; + } + } + } + + const ReturnStatus group_end_status = Api::Check(Api::GroupEnd()); + return first_error != ReturnStatus::kSuccess ? first_error + : group_end_status; + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_GATHER_H_ diff --git a/src/backends/ccl/mccl/api.h b/src/backends/ccl/mccl/api.h index a5d2bcd..2e97324 100644 --- a/src/backends/ccl/mccl/api.h +++ b/src/backends/ccl/mccl/api.h @@ -49,6 +49,20 @@ struct McclApi { return mcclAllReduce(send_buff, recv_buff, count, data_type, op, comm, stream); } + + static Result GroupStart() { return mcclGroupStart(); } + + static Result GroupEnd() { return mcclGroupEnd(); } + + static Result Send(const void *send_buff, size_t count, DataType data_type, + int peer, Comm comm, Stream stream) { + return mcclSend(send_buff, count, data_type, peer, comm, stream); + } + + static Result Recv(void *recv_buff, size_t count, DataType data_type, + int peer, Comm comm, Stream stream) { + return mcclRecv(recv_buff, count, data_type, peer, comm, stream); + } }; } // namespace infini::ccl diff --git a/src/backends/ccl/mccl/impl/gather.h b/src/backends/ccl/mccl/impl/gather.h new file mode 100644 index 0000000..8c2f158 --- /dev/null +++ b/src/backends/ccl/mccl/impl/gather.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_GATHER_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_GATHER_H_ + +#include "backends/ccl/common/impl/gather.h" + +namespace infini::ccl { + +template +class GatherImpl + : public CclGatherImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_GATHER_H_ diff --git a/src/backends/ccl/nccl/api.h b/src/backends/ccl/nccl/api.h index e7b6119..8dd49a2 100644 --- a/src/backends/ccl/nccl/api.h +++ b/src/backends/ccl/nccl/api.h @@ -46,6 +46,20 @@ struct NcclApi { return ncclAllReduce(send_buff, recv_buff, count, data_type, op, comm, stream); } + + static Result GroupStart() { return ncclGroupStart(); } + + static Result GroupEnd() { return ncclGroupEnd(); } + + 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); + } + + static Result Recv(void *recv_buff, size_t count, DataType data_type, + int peer, Comm comm, Stream stream) { + return ncclRecv(recv_buff, count, data_type, peer, comm, stream); + } }; } // namespace infini::ccl diff --git a/src/backends/ccl/nccl/impl/gather.h b/src/backends/ccl/nccl/impl/gather.h new file mode 100644 index 0000000..28ed47d --- /dev/null +++ b/src/backends/ccl/nccl/impl/gather.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_GATHER_H_ +#define INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_GATHER_H_ + +#include "backends/ccl/common/impl/gather.h" + +namespace infini::ccl { + +template +class GatherImpl + : public CclGatherImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_GATHER_H_ diff --git a/src/backends/mpi/ompi/impl/gather.h b/src/backends/mpi/ompi/impl/gather.h index ee7dd1f..b88806e 100644 --- a/src/backends/mpi/ompi/impl/gather.h +++ b/src/backends/mpi/ompi/impl/gather.h @@ -3,6 +3,7 @@ #include #include +#include #include "backends/mpi/ompi/checks.h" #include "backends/mpi/ompi/comm_instance.h" @@ -24,11 +25,20 @@ class GatherImpl { ListGetBest(ActiveDevices{}); using Rt = Runtime; - auto *inst = static_cast(comm->inter_comm()); - if (!inst || inst->handle == MPI_COMM_NULL) { + if (!comm || !comm->inter_comm() || + comm->inter_comm_backend() != BackendType::kOmpi) { LOG("Invalid OpenMPI communicator instance for `Gather`."); return ReturnStatus::kInternalError; } + auto *inst = static_cast(comm->inter_comm()); + if (inst->handle == MPI_COMM_NULL) { + LOG("Invalid OpenMPI communicator handle for `Gather`."); + return ReturnStatus::kInternalError; + } + if (comm->size() <= 0 || comm->rank() < 0 || comm->rank() >= comm->size()) { + LOG("Invalid rank or world size for `Gather`."); + return ReturnStatus::kInternalError; + } size_t type_size = kDataTypeToSize.at(data_type); if (count > std::numeric_limits::max() / type_size) { @@ -36,45 +46,44 @@ class GatherImpl { return ReturnStatus::kInvalidArgument; } size_t send_bytes = count * type_size; - size_t recv_bytes = send_bytes * static_cast(comm->size()); - const bool is_root = comm->rank() == root; - - // Transfer raw bytes so the gather is correct for every data type, - // including `kFloat16` / `kBFloat16`, which map to `MPI_BYTE`. if (send_bytes > static_cast(std::numeric_limits::max())) { LOG("Per-rank byte count exceeds MPI int range for `Gather`."); return ReturnStatus::kInvalidArgument; } + const size_t world_size = static_cast(comm->size()); + if (send_bytes > std::numeric_limits::max() / world_size) { + LOG("Total byte size overflows `size_t` for `Gather`."); + return ReturnStatus::kInvalidArgument; + } + const size_t recv_bytes = send_bytes * world_size; + const bool is_root = comm->rank() == root; int mpi_byte_count = static_cast(send_bytes); - // Host staging buffers. Only `root` allocates the receive side, since - // `MPI_Gather` writes the gathered result only on `root`. - void *host_sendbuf = std::malloc(send_bytes); - void *host_recvbuf = is_root ? std::malloc(recv_bytes) : nullptr; + // Transfer raw bytes so movement-only collectives preserve every InfiniCCL + // data type, including float16 and bfloat16. + std::unique_ptr host_sendbuf( + std::malloc(send_bytes), &std::free); + std::unique_ptr host_recvbuf( + is_root ? std::malloc(recv_bytes) : nullptr, &std::free); if (!host_sendbuf || (is_root && !host_recvbuf)) { - std::free(host_sendbuf); - std::free(host_recvbuf); LOG("Failed to allocate host buffers for `Gather` staging."); return ReturnStatus::kSystemError; } - CHECK_STATUS(Rt, Rt::Memcpy(host_sendbuf, send_buff, send_bytes, + CHECK_STATUS(Rt, Rt::Memcpy(host_sendbuf.get(), send_buff, send_bytes, Rt::MemcpyDeviceToHost)); CHECK_STATUS(Rt, Rt::StreamSynchronize(static_cast(stream))); // Note: `MPI_Gather`'s `recvcount` is the per-rank count, not the total. - INFINI_CHECK_MPI(MPI_Gather(host_sendbuf, mpi_byte_count, MPI_BYTE, - host_recvbuf, mpi_byte_count, MPI_BYTE, root, - inst->handle)); + INFINI_CHECK_MPI(MPI_Gather(host_sendbuf.get(), mpi_byte_count, MPI_BYTE, + host_recvbuf.get(), mpi_byte_count, MPI_BYTE, + root, inst->handle)); if (is_root) { - CHECK_STATUS(Rt, Rt::Memcpy(recv_buff, host_recvbuf, recv_bytes, + CHECK_STATUS(Rt, Rt::Memcpy(recv_buff, host_recvbuf.get(), recv_bytes, Rt::MemcpyHostToDevice)); } - std::free(host_sendbuf); - std::free(host_recvbuf); - return ReturnStatus::kSuccess; } }; diff --git a/src/base/gather.h b/src/base/gather.h index ec9c438..5ece773 100644 --- a/src/base/gather.h +++ b/src/base/gather.h @@ -18,17 +18,20 @@ class Gather : public Operation { static ReturnStatus Execute(const void *send_buff, void *recv_buff, size_t count, DataType datatype, int root, void *comm_handle, void *stream) { - if (!comm_handle) { - LOG("Invalid communicator handle for `Gather`."); + if (HasInvalidRequiredArgs(datatype, root, comm_handle)) { return ReturnStatus::kInvalidArgument; } - + if (count == 0) { + return ReturnStatus::kSuccess; + } auto *comm = static_cast(comm_handle); - if (HasInvalidArgs(send_buff, recv_buff, datatype, root, comm)) { + if (!send_buff) { + LOG("Invalid send buffer pointer for `Gather`."); return ReturnStatus::kInvalidArgument; } - if (count == 0) { - return ReturnStatus::kSuccess; + if (comm->rank() == root && !recv_buff) { + LOG("Invalid root receive buffer pointer for `Gather`."); + return ReturnStatus::kInvalidArgument; } return GatherImpl::Apply( @@ -36,24 +39,21 @@ class Gather : public Operation { } private: - static bool HasInvalidArgs(const void *send_buff, void *recv_buff, - DataType datatype, int root, Communicator *comm) { + static bool HasInvalidRequiredArgs(DataType datatype, int root, + void *comm_handle) { + if (!comm_handle) { + LOG("Invalid communicator handle for `Gather`."); + return true; + } if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) { LOG("Invalid data type for `Gather`."); return true; } + auto *comm = static_cast(comm_handle); if (root < 0 || root >= comm->size()) { LOG("Invalid root rank for `Gather`."); return true; } - if (!send_buff) { - LOG("Invalid send buffer pointer for `Gather`."); - return true; - } - if (comm->rank() == root && !recv_buff) { - LOG("Invalid root receive buffer pointer for `Gather`."); - return true; - } return false; } };