From baadfba338f4564e0da038c6d74c96033df0e965 Mon Sep 17 00:00:00 2001 From: GordonYang1 <1468121796@qq.com> Date: Sat, 22 Aug 2026 22:48:50 +0800 Subject: [PATCH] feat: support CCL `Scatter` --- examples/ccl/scatter.cc | 332 ++++++++++++++++++++++++ examples/ccl_mpi_hybrid/scatter.cc | 345 +++++++++++++++++++++++++ examples/mpi/scatter.cc | 296 +++++++++++++++------ src/backends/ccl/common/impl/scatter.h | 113 ++++++++ src/backends/ccl/mccl/api.h | 14 + src/backends/ccl/mccl/impl/scatter.h | 17 ++ src/backends/ccl/nccl/api.h | 14 + src/backends/ccl/nccl/impl/scatter.h | 17 ++ src/backends/mpi/ompi/impl/scatter.h | 52 ++-- src/base/scatter.h | 32 +-- 10 files changed, 1113 insertions(+), 119 deletions(-) create mode 100644 examples/ccl/scatter.cc create mode 100644 examples/ccl_mpi_hybrid/scatter.cc create mode 100644 src/backends/ccl/common/impl/scatter.h create mode 100644 src/backends/ccl/mccl/impl/scatter.h create mode 100644 src/backends/ccl/nccl/impl/scatter.h diff --git a/examples/ccl/scatter.cc b/examples/ccl/scatter.cc new file mode 100644 index 0000000..95d4e62 --- /dev/null +++ b/examples/ccl/scatter.cc @@ -0,0 +1,332 @@ +/** + * InfiniCCL Example: Thread-per-GPU Single-Node Scatter + * + * This example creates one native CCL rank per GPU and scatters one distinct + * block from rank 0 to every rank 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}; + std::vector samples; + + explicit ScenarioState(int world_size) + : samples(static_cast(world_size), 0.0f) {} +}; + +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; +} + +void FillScatterInput(std::vector *input, size_t num_elements, + int world_size) { + for (int destination = 0; destination < world_size; ++destination) { + const size_t offset = static_cast(destination) * num_elements; + std::fill_n(input->begin() + offset, num_elements, + static_cast(destination + 1)); + } +} + +void PrintScatterMetrics(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 total_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 (" + << total_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 = + total_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 &samples, + 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 Scatter 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 rank = 0; rank < std::min(world_size, 4); ++rank) { + std::cout << "[r" << rank << ": " << samples[static_cast(rank)] + << "] "; + } + std::cout << std::endl; + PrintScatterMetrics(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 Scatter 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 total_elements = + args.num_elements * static_cast(args.size); + const size_t total_bytes = total_elements * sizeof(float); + std::vector h_send; + if (args.rank == kRoot) { + h_send.resize(total_elements, 0.0f); + FillScatterInput(&h_send, args.num_elements, args.size); + } + std::vector h_recv(args.num_elements, 0.0f); + + float *d_send = nullptr; + float *d_recv = nullptr; + if (args.rank == kRoot) { + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), total_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes, + Rt::MemcpyHostToDevice)); + } + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_recv), rank_bytes)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + for (int i = 0; i < args.warmup_iterations; ++i) { + CHECK_INFINI(infinicclScatter(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(infinicclScatter(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); + + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, rank_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const bool local_correct = + Validator::ValidateResult(h_recv.data(), args.num_elements, + static_cast(args.rank + 1), args.rank); + args.state->samples[static_cast(args.rank)] = h_recv.front(); + if (!local_correct) { + args.state->correct.store(false, std::memory_order_release); + } + + WaitForAll(args.state, args.size); + if (args.rank == kRoot) { + PrintResult(args.state->correct.load(std::memory_order_acquire), + args.state->samples, args.num_elements, args.size, elapsed_ms); + } + + if (args.rank == kRoot) { + 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 sent to 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 Scatter." << std::endl; + return EXIT_FAILURE; + } + } + + if (optind != argc) { + std::cerr << "Unexpected positional argument for Scatter." << 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 << "Scatter 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 Scatter." << 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(num_gpus); + 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 Scatter validation passed." << std::endl; + } else { + std::cerr << "[Main Process] CCL Scatter 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/scatter.cc b/examples/ccl_mpi_hybrid/scatter.cc new file mode 100644 index 0000000..6868b28 --- /dev/null +++ b/examples/ccl_mpi_hybrid/scatter.cc @@ -0,0 +1,345 @@ +/** + * InfiniCCL Example: Scatter (OpenMPI + CCL Hybrid) + * + * This example first exercises Scatter 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; +} + +void FillScatterInput(std::vector *input, size_t num_elements, + int world_size) { + for (int destination = 0; destination < world_size; ++destination) { + const size_t offset = static_cast(destination) * num_elements; + std::fill_n(input->begin() + offset, num_elements, + static_cast(destination + 1)); + } +} + +template +bool CollectValidationReports(bool local_correct, float local_sample, int rank, + int world_size, infinicclComm_t comm, + std::vector *root_reports) { + const float report = local_correct ? local_sample : -1.0f; + float *d_report = nullptr; + float *d_reports = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_report), sizeof(float))); + if (rank == kRoot) { + root_reports->assign(static_cast(world_size), 0.0f); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_reports), + static_cast(world_size) * sizeof(float))); + } + CHECK_RT( + Rt, Rt::Memcpy(d_report, &report, sizeof(float), Rt::MemcpyHostToDevice)); + CHECK_INFINI(infinicclGather(d_report, d_reports, 1, infinicclFloat32, kRoot, + comm, nullptr)); + + int32_t global_status = 1; + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Memcpy(root_reports->data(), d_reports, + static_cast(world_size) * sizeof(float), + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + for (int source = 0; source < world_size; ++source) { + if ((*root_reports)[static_cast(source)] != + static_cast(source + 1)) { + global_status = 0; + } + } + } + CHECK_INFINI(infinicclBroadcast(&global_status, &global_status, 1, + infinicclInt32, kRoot, comm, nullptr)); + + CHECK_RT(Rt, Rt::Free(d_report)); + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Free(d_reports)); + } + return global_status == 1; +} + +void PrintScatterMetrics(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 total_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 (" + << total_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 = + total_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 &samples, + 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 Scatter 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 rank = 0; rank < std::min(world_size, 4); ++rank) { + std::cout << "[r" << rank << ": " << samples[static_cast(rank)] + << "] "; + } + std::cout << std::endl; + PrintScatterMetrics(num_elements, world_size, elapsed_ms); +} + +bool RunScatterExample(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 Scatter." << 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 Scatter." + << 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 Scatter provider + // delegates this rank token distribution to the OpenMPI inter communicator. + std::vector h_bootstrap_send; + if (rank == kRoot) { + h_bootstrap_send.resize(static_cast(size)); + for (int destination = 0; destination < size; ++destination) { + h_bootstrap_send[static_cast(destination)] = + static_cast(destination + 1); + } + } + float h_bootstrap_recv = 0.0f; + float *d_bootstrap_send = nullptr; + float *d_bootstrap_recv = nullptr; + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_bootstrap_send), + static_cast(size) * sizeof(float))); + CHECK_RT(Rt, Rt::Memcpy(d_bootstrap_send, h_bootstrap_send.data(), + static_cast(size) * sizeof(float), + Rt::MemcpyHostToDevice)); + } + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_bootstrap_recv), + sizeof(float))); + CHECK_INFINI(infinicclScatter(d_bootstrap_send, d_bootstrap_recv, 1, + infinicclFloat32, kRoot, comm, nullptr)); + CHECK_RT(Rt, Rt::Memcpy(&h_bootstrap_recv, d_bootstrap_recv, sizeof(float), + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + const bool bootstrap_local_correct = h_bootstrap_recv == rank + 1.0f; + std::vector bootstrap_reports; + const bool bootstrap_correct = + CollectValidationReports(bootstrap_local_correct, h_bootstrap_recv, + rank, size, comm, &bootstrap_reports); + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Free(d_bootstrap_send)); + } + CHECK_RT(Rt, Rt::Free(d_bootstrap_recv)); + + if (!bootstrap_correct) { + if (rank == kRoot) { + std::cerr << "OpenMPI Scatter fallback validation failed." << std::endl; + } + CHECK_INFINI(infinicclCommDestroy(comm)); + CHECK_INFINI(infinicclFinalize()); + return false; + } + if (rank == kRoot) { + std::cout << "OpenMPI Scatter 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 Scatter buffer size overflows `size_t`." << std::endl; + std::exit(EXIT_FAILURE); + } + const size_t rank_bytes = kNumElements * sizeof(float); + const size_t total_elements = kNumElements * world_size; + const size_t total_bytes = total_elements * sizeof(float); + std::vector h_send; + if (rank == kRoot) { + h_send.resize(total_elements, 0.0f); + FillScatterInput(&h_send, kNumElements, size); + } + std::vector h_recv(kNumElements, 0.0f); + + float *d_send = nullptr; + float *d_recv = nullptr; + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), total_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes, + Rt::MemcpyHostToDevice)); + } + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_recv), rank_bytes)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + for (int i = 0; i < kWarmupIterations; ++i) { + CHECK_INFINI(infinicclScatter(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(infinicclScatter(d_send, d_recv, kNumElements, + infinicclFloat32, kRoot, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const double elapsed_ms = + timer.ElapsedMs() / static_cast(kProfileIterations); + + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, rank_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const bool local_correct = Validator::ValidateResult( + h_recv.data(), kNumElements, static_cast(rank + 1), rank); + std::vector validation_reports; + const bool correct = CollectValidationReports( + local_correct, h_recv.front(), rank, size, comm, &validation_reports); + + if (rank == kRoot) { + PrintResult(correct, validation_reports, kNumElements, size, elapsed_ms); + } + + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Free(d_send)); + } + 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 Scatter validation passed." + << std::endl; + } else { + std::cerr << "[Main Process] Hybrid CCL Scatter validation failed." + << std::endl; + } + std::cout << "InfiniCCL finalized." << std::endl; + } + return correct; +} + +} // namespace + +int main(int argc, char **argv) { + return RunScatterExample(argc, argv) ? EXIT_SUCCESS : EXIT_FAILURE; +} diff --git a/examples/mpi/scatter.cc b/examples/mpi/scatter.cc index 2c16577..4363b61 100644 --- a/examples/mpi/scatter.cc +++ b/examples/mpi/scatter.cc @@ -1,13 +1,24 @@ /** - * InfiniCCL Example: Scatter - * * This example demonstrates the API for performing a collective - * data distribution across multiple GPUs and nodes, where `root` - * sends a distinct block to every rank. + * InfiniCCL Example: Scatter (MPI Backend) + * + * Rank 0 distributes one GPU-resident block to every rank. Each rank validates + * its full block, and rank 0 reports Scatter-specific bandwidth. */ #include +#include +#include +#include +#include +#include +#include +#include +#include #include +#include +#include +#include #include // Public API @@ -25,129 +36,250 @@ using namespace infini::ccl; -void RunScatterExample(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; +} + +void FillScatterInput(std::vector *input, size_t num_elements, + int world_size) { + for (int destination = 0; destination < world_size; ++destination) { + const size_t offset = static_cast(destination) * num_elements; + std::fill_n(input->begin() + offset, num_elements, + static_cast(destination + 1)); + } +} + +template +bool CollectValidationReports(bool local_correct, float local_sample, int rank, + int world_size, infinicclComm_t comm, + std::vector *root_reports) { + const float report = local_correct ? local_sample : -1.0f; + float *d_report = nullptr; + float *d_reports = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_report), sizeof(float))); + if (rank == kRoot) { + root_reports->assign(static_cast(world_size), 0.0f); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_reports), + static_cast(world_size) * sizeof(float))); + } + CHECK_RT( + Rt, Rt::Memcpy(d_report, &report, sizeof(float), Rt::MemcpyHostToDevice)); + CHECK_INFINI(infinicclGather(d_report, d_reports, 1, infinicclFloat32, kRoot, + comm, nullptr)); + + int32_t global_status = 1; + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Memcpy(root_reports->data(), d_reports, + static_cast(world_size) * sizeof(float), + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + for (int source = 0; source < world_size; ++source) { + if ((*root_reports)[static_cast(source)] != + static_cast(source + 1)) { + global_status = 0; + } + } + } + CHECK_INFINI(infinicclBroadcast(&global_status, &global_status, 1, + infinicclInt32, kRoot, comm, nullptr)); + + CHECK_RT(Rt, Rt::Free(d_report)); + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Free(d_reports)); + } + return global_status == 1; +} + +void PrintScatterMetrics(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 total_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 (" + << total_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 = + total_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 RunScatterExample(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 Scatter." << 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 Scatter." << 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 Scatter - constexpr int kRoot = 0; - - // Prepare Data - std::vector h_send(kNumElements * size, 0.0f); - std::vector h_recv(kNumElements, 0.0f); - - // Initialize: `root` fills the block destined for rank `r` with `(r + 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 Scatter buffer size overflows `size_t`." << std::endl; + std::exit(EXIT_FAILURE); + } + const size_t rank_bytes = num_elements * sizeof(float); + const size_t total_elements = num_elements * world_size; + const size_t total_bytes = total_elements * sizeof(float); + std::vector h_send; if (rank == kRoot) { - for (int r = 0; r < size; ++r) { - for (size_t i = 0; i < kNumElements; ++i) { - h_send[static_cast(r) * kNumElements + i] = - static_cast(r + 1); - } - } + h_send.resize(total_elements, 0.0f); + FillScatterInput(&h_send, num_elements, size); } + std::vector h_recv(num_elements, 0.0f); - float *d_send, *d_recv; - size_t recv_bytes = kNumElements * sizeof(*d_recv); - size_t send_bytes = recv_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)); - + float *d_send = nullptr; + float *d_recv = nullptr; if (rank == kRoot) { - std::cout << "\n=== Performing Scatter on GPU Memory ===" << std::endl; - std::cout << "Data size per rank: " << kNumElements << " floats (" - << recv_bytes / 1024 / 1024 << " MB)" << std::endl; - std::cout << "Operation: Scatter" << 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; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), total_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes, + Rt::MemcpyHostToDevice)); } - + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_recv), rank_bytes)); CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); - // Warm-up - CHECK_INFINI(infinicclScatter(d_send, d_recv, kNumElements, infinicclFloat32, - kRoot, comm, nullptr)); + if (rank == kRoot) { + std::cout << "\n=== Performing MPI Scatter 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(infinicclScatter(d_send, d_recv, kNumElements, + for (int i = 0; i < warmup_iterations; ++i) { + CHECK_INFINI(infinicclScatter(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(infinicclScatter(d_send, d_recv, kNumElements, + for (int i = 0; i < profile_iterations; ++i) { + CHECK_INFINI(infinicclScatter(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: every rank should receive its own `(rank + 1)` block. - CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, recv_bytes, + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, rank_bytes, Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const bool local_correct = Validator::ValidateResult( + h_recv.data(), num_elements, static_cast(rank + 1), rank); + std::vector validation_reports; + const bool correct = CollectValidationReports( + local_correct, h_recv.front(), rank, size, comm, &validation_reports); - Validator::ValidateResult(h_recv.data(), kNumElements, - static_cast(rank + 1), rank, true, - "Scatter"); - - // Metrics Reporting (Only from rank 0 for cleaner output) if (rank == kRoot) { - Metrics metrics{elapsed, send_bytes, size}; - metrics.Print(); + constexpr const char *kGreen = "\033[32m"; + constexpr const char *kRed = "\033[31m"; + constexpr const char *kReset = "\033[0m"; + std::cout << "\n=== MPI Scatter 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(size, 4); ++source) { + std::cout << "[r" << source << ": " + << validation_reports[static_cast(source)] << "] "; + } + std::cout << std::endl; + PrintScatterMetrics(num_elements, size, elapsed_ms); } - // Cleanup - CHECK_RT(Rt, Rt::Free(d_send)); + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Free(d_send)); + } 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; - - RunScatterExample(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 RunScatterExample(argc, argv, kWarmupIterations, kProfileIterations, + kNumElements) + ? EXIT_SUCCESS + : EXIT_FAILURE; } diff --git a/src/backends/ccl/common/impl/scatter.h b/src/backends/ccl/common/impl/scatter.h new file mode 100644 index 0000000..3674861 --- /dev/null +++ b/src/backends/ccl/common/impl/scatter.h @@ -0,0 +1,113 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_SCATTER_H_ +#define INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_SCATTER_H_ + +#include +#include + +#include "backends/ccl/common/api.h" +#include "backends/ccl/common/comm_instance.h" +#include "base/scatter.h" +#include "communicator.h" +#include "data_type_impl.h" +#include "logging.h" + +namespace infini::ccl { + +template +struct DeferredScatter { + using type = Scatter; +}; + +template +class CclScatterImpl { + 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 DeferredScatter::type; + if constexpr (BackendEnabled::value) { + return ScatterImpl::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 `Scatter`."); + 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 `Scatter`."); + 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 `Scatter`."); + return ReturnStatus::kInvalidArgument; + } + + auto native_stream = reinterpret_cast(stream); + ReturnStatus status = Api::Check(Api::GroupStart()); + if (status != ReturnStatus::kSuccess) { + return status; + } + + ReturnStatus first_error = ReturnStatus::kSuccess; + if (comm->rank() == root) { + const auto *send_bytes = static_cast(send_buff); + for (int peer = 0; peer < comm->size(); ++peer) { + const size_t offset = static_cast(peer) * rank_bytes; + status = Api::Check(Api::Send(send_bytes + offset, count, native_type, + peer, instance->handle, native_stream)); + if (first_error == ReturnStatus::kSuccess && + status != ReturnStatus::kSuccess) { + first_error = status; + } + } + } + + status = Api::Check(Api::Recv(recv_buff, count, native_type, root, + 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_SCATTER_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/scatter.h b/src/backends/ccl/mccl/impl/scatter.h new file mode 100644 index 0000000..438ad46 --- /dev/null +++ b/src/backends/ccl/mccl/impl/scatter.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_SCATTER_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_SCATTER_H_ + +#include "backends/ccl/common/impl/scatter.h" + +namespace infini::ccl { + +template +class ScatterImpl + : public CclScatterImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_SCATTER_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/scatter.h b/src/backends/ccl/nccl/impl/scatter.h new file mode 100644 index 0000000..b710f00 --- /dev/null +++ b/src/backends/ccl/nccl/impl/scatter.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_SCATTER_H_ +#define INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_SCATTER_H_ + +#include "backends/ccl/common/impl/scatter.h" + +namespace infini::ccl { + +template +class ScatterImpl + : public CclScatterImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_SCATTER_H_ diff --git a/src/backends/mpi/ompi/impl/scatter.h b/src/backends/mpi/ompi/impl/scatter.h index 445c63c..2a4e5c6 100644 --- a/src/backends/mpi/ompi/impl/scatter.h +++ b/src/backends/mpi/ompi/impl/scatter.h @@ -3,6 +3,7 @@ #include #include +#include #include "backends/mpi/ompi/checks.h" #include "backends/mpi/ompi/comm_instance.h" @@ -24,11 +25,21 @@ class ScatterImpl { 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 `Scatter`."); return ReturnStatus::kInternalError; } + auto *inst = static_cast(comm->inter_comm()); + if (inst->handle == MPI_COMM_NULL) { + LOG("Invalid OpenMPI communicator handle for `Scatter`."); + 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 `Scatter`."); + return ReturnStatus::kInternalError; + } size_t type_size = kDataTypeToSize.at(data_type); if (count > std::numeric_limits::max() / type_size) { @@ -36,45 +47,44 @@ class ScatterImpl { return ReturnStatus::kInvalidArgument; } size_t recv_bytes = count * type_size; - size_t send_bytes = recv_bytes * static_cast(comm->size()); - const bool is_root = comm->rank() == root; - - // Transfer raw bytes so the scatter is correct for every data type, - // including `kFloat16` / `kBFloat16`, which map to `MPI_BYTE`. if (recv_bytes > static_cast(std::numeric_limits::max())) { LOG("Per-rank byte count exceeds MPI int range for `Scatter`."); return ReturnStatus::kInvalidArgument; } + const size_t world_size = static_cast(comm->size()); + if (recv_bytes > std::numeric_limits::max() / world_size) { + LOG("Total byte size overflows `size_t` for `Scatter`."); + return ReturnStatus::kInvalidArgument; + } + const size_t send_bytes = recv_bytes * world_size; + const bool is_root = comm->rank() == root; int mpi_byte_count = static_cast(recv_bytes); - // Host staging buffers. Only `root` allocates the send side, since - // `MPI_Scatter` reads the distributed blocks only from `root`. - void *host_sendbuf = is_root ? std::malloc(send_bytes) : nullptr; - void *host_recvbuf = std::malloc(recv_bytes); + // Transfer raw bytes so movement-only collectives preserve every InfiniCCL + // data type, including float16 and bfloat16. + std::unique_ptr host_sendbuf( + is_root ? std::malloc(send_bytes) : nullptr, &std::free); + std::unique_ptr host_recvbuf( + std::malloc(recv_bytes), &std::free); if ((is_root && !host_sendbuf) || !host_recvbuf) { - std::free(host_sendbuf); - std::free(host_recvbuf); LOG("Failed to allocate host buffers for `Scatter` staging."); return ReturnStatus::kSystemError; } if (is_root) { - 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_Scatter`'s `sendcount` is the per-rank count, not the total. - INFINI_CHECK_MPI(MPI_Scatter(host_sendbuf, mpi_byte_count, MPI_BYTE, - host_recvbuf, mpi_byte_count, MPI_BYTE, root, - inst->handle)); + INFINI_CHECK_MPI(MPI_Scatter(host_sendbuf.get(), mpi_byte_count, MPI_BYTE, + host_recvbuf.get(), mpi_byte_count, MPI_BYTE, + root, inst->handle)); - 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/scatter.h b/src/base/scatter.h index 62643f6..4299622 100644 --- a/src/base/scatter.h +++ b/src/base/scatter.h @@ -18,17 +18,20 @@ class Scatter : 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 `Scatter`."); + 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 (!recv_buff) { + LOG("Invalid receive buffer pointer for `Scatter`."); return ReturnStatus::kInvalidArgument; } - if (count == 0) { - return ReturnStatus::kSuccess; + if (comm->rank() == root && !send_buff) { + LOG("Invalid root send buffer pointer for `Scatter`."); + return ReturnStatus::kInvalidArgument; } return ScatterImpl::Apply( @@ -36,24 +39,21 @@ class Scatter : 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 `Scatter`."); + return true; + } if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) { LOG("Invalid data type for `Scatter`."); return true; } + auto *comm = static_cast(comm_handle); if (root < 0 || root >= comm->size()) { LOG("Invalid root rank for `Scatter`."); return true; } - if (!recv_buff) { - LOG("Invalid receive buffer pointer for `Scatter`."); - return true; - } - if (comm->rank() == root && !send_buff) { - LOG("Invalid root send buffer pointer for `Scatter`."); - return true; - } return false; } };