diff --git a/examples/ccl/reduce_scatter.cc b/examples/ccl/reduce_scatter.cc new file mode 100644 index 0000000..40f560e --- /dev/null +++ b/examples/ccl/reduce_scatter.cc @@ -0,0 +1,322 @@ +/** + * InfiniCCL Example: Thread-per-GPU Single-Node ReduceScatter + * + * This example creates one native CCL rank per GPU, then validates + * out-of-place and canonical in-place ReduceScatter 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 recv_count; + 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 FillInput(std::vector *input, size_t recv_count, int world_size, + int rank) { + for (int dest_rank = 0; dest_rank < world_size; ++dest_rank) { + const float value = + static_cast(rank + 1) * static_cast(dest_rank + 1); + const size_t offset = static_cast(dest_rank) * recv_count; + std::fill_n(input->begin() + offset, recv_count, value); + } +} + +float ExpectedValue(int rank, int world_size) { + const float rank_sum = static_cast(world_size) * + (static_cast(world_size) + 1.0f) / 2.0f; + return rank_sum * static_cast(rank + 1); +} + +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(); + } +} + +void PrintResult(const char *scenario, bool correct, + const std::vector &result, float expected, + 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 ReduceScatter Results ===" << std::endl; + std::cout << "Correct: " + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) + << std::endl; + std::cout << "Expect: " << expected << std::endl; + std::cout << "Actual: " << result.front() << std::endl; + std::cout << "Average time: " << elapsed_ms << " ms" << std::endl; +} + +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 ReduceScatter 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_count = args.recv_count * static_cast(args.size); + const size_t send_bytes = send_count * sizeof(float); + const size_t recv_bytes = args.recv_count * sizeof(float); + const size_t local_offset = static_cast(args.rank) * args.recv_count; + + std::vector h_send(send_count); + std::vector h_recv(args.recv_count, 0.0f); + FillInput(&h_send, args.recv_count, args.size, args.rank); + + 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)); + + auto run_scenario = [&](bool in_place) { + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), send_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + float *active_recv = in_place ? d_send + local_offset : d_recv; + auto restore_local_block = [&]() { + if (in_place) { + CHECK_RT(Rt, Rt::Memcpy(active_recv, h_send.data() + local_offset, + recv_bytes, Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + } + }; + + for (int i = 0; i < args.warmup_iter; ++i) { + restore_local_block(); + CHECK_INFINI(infinicclReduceScatter(d_send, active_recv, args.recv_count, + infinicclFloat32, infinicclSum, comm, + nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + double elapsed_ms = 0.0; + if (in_place) { + for (int i = 0; i < args.profile_iter; ++i) { + restore_local_block(); + Timer timer; + CHECK_INFINI(infinicclReduceScatter(d_send, active_recv, + args.recv_count, infinicclFloat32, + infinicclSum, comm, nullptr)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + elapsed_ms += timer.ElapsedMs(); + } + elapsed_ms /= static_cast(args.profile_iter); + } else { + Timer timer; + for (int i = 0; i < args.profile_iter; ++i) { + CHECK_INFINI(infinicclReduceScatter(d_send, active_recv, + args.recv_count, infinicclFloat32, + infinicclSum, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + elapsed_ms = timer.ElapsedMs() / static_cast(args.profile_iter); + } + + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), active_recv, recv_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + const float expected = ExpectedValue(args.rank, args.size); + const bool correct = Validator::ValidateResult( + h_recv.data(), args.recv_count, expected, 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, + expected, 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 received 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 recv_count = 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, &recv_count); + 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 ReduceScatter." + << std::endl; + return EXIT_FAILURE; + } + } + + if (optind != argc) { + std::cerr << "Unexpected positional argument for ReduceScatter." + << std::endl; + return EXIT_FAILURE; + } + if (recv_count > std::numeric_limits::max() / sizeof(float) || + static_cast(num_gpus) > + std::numeric_limits::max() / recv_count || + recv_count * static_cast(num_gpus) > + std::numeric_limits::max() / sizeof(float)) { + std::cerr << "ReduceScatter 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 ReduceScatter." << 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, recv_count, + 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 ReduceScatter scenarios passed." + << std::endl; + } else { + std::cerr << "[Main Process] ReduceScatter 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/reduce_scatter.cc b/examples/ccl_mpi_hybrid/reduce_scatter.cc new file mode 100644 index 0000000..69ce9cf --- /dev/null +++ b/examples/ccl_mpi_hybrid/reduce_scatter.cc @@ -0,0 +1,320 @@ +/** + * InfiniCCL Example: ReduceScatter (OpenMPI + CCL Hybrid) + * + * This example first uses ReduceScatter 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 canonical in-place GPU + * ReduceScatter. + */ + +#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; +} + +void FillInput(std::vector *input, size_t recv_count, int world_size, + int rank) { + for (int dest_rank = 0; dest_rank < world_size; ++dest_rank) { + const float value = + static_cast(rank + 1) * static_cast(dest_rank + 1); + const size_t offset = static_cast(dest_rank) * recv_count; + std::fill_n(input->begin() + offset, recv_count, value); + } +} + +float ExpectedValue(int rank, int world_size) { + const float rank_sum = static_cast(world_size) * + (static_cast(world_size) + 1.0f) / 2.0f; + return rank_sum * static_cast(rank + 1); +} + +void PrintResult(const char *scenario, bool correct, float actual, + float expected, 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 ReduceScatter Results ===" << std::endl; + std::cout << "Correct: " + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) + << std::endl; + std::cout << "Expect: " << expected << std::endl; + std::cout << "Actual: " << actual << std::endl; + std::cout << "Average time: " << elapsed_ms << " ms" << std::endl; +} + +bool RunReduceScatterExample(int argc, char **argv) { + constexpr Device::Type kDevType = + ListGetBest(EnabledDevices{}); + using Rt = Runtime; + + constexpr int kWarmupIterations = 2; + constexpr int kProfileIterations = 20; + constexpr size_t kRecvCount = 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 ReduceScatter." << 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 ReduceScatter." + << 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 ReduceScatter itself. Rank 0 + // repeats the same unique ID in every destination block; all other ranks + // contribute zeros. With only the OpenMPI inter communicator initialized, + // the CCL provider delegates this call to MPI_Reduce_scatter_block. + infinicclUniqueId id{}; + if (rank == 0) { + CHECK_INFINI(infinicclGetUniqueId(&id)); + } + const size_t id_bytes = sizeof(id); + const size_t world_size = static_cast(size); + if (world_size > std::numeric_limits::max() / id_bytes) { + std::cerr << "ReduceScatter bootstrap buffer size overflows `size_t`." + << std::endl; + std::exit(EXIT_FAILURE); + } + const size_t id_send_bytes = id_bytes * world_size; + std::vector h_id_send(id_send_bytes, 0); + if (rank == 0) { + for (int dest_rank = 0; dest_rank < size; ++dest_rank) { + std::memcpy(h_id_send.data() + static_cast(dest_rank) * id_bytes, + &id, id_bytes); + } + } + + uint8_t *d_id_send = nullptr; + uint8_t *d_id_recv = nullptr; + CHECK_RT(Rt, + Rt::Malloc(reinterpret_cast(&d_id_send), id_send_bytes)); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_id_recv), id_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_id_send, h_id_send.data(), id_send_bytes, + Rt::MemcpyHostToDevice)); + CHECK_INFINI(infinicclReduceScatter(d_id_send, d_id_recv, id_bytes, + infinicclUInt8, infinicclSum, 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 (kRecvCount > std::numeric_limits::max() / world_size || + kRecvCount * world_size > + std::numeric_limits::max() / sizeof(float)) { + std::cerr << "Hybrid ReduceScatter buffer size overflows `size_t`." + << std::endl; + std::exit(EXIT_FAILURE); + } + + const size_t send_count = kRecvCount * world_size; + const size_t send_bytes = send_count * sizeof(float); + const size_t recv_bytes = kRecvCount * sizeof(float); + const size_t local_offset = static_cast(rank) * kRecvCount; + std::vector h_send(send_count); + std::vector h_recv(kRecvCount, 0.0f); + FillInput(&h_send, kRecvCount, size, rank); + + 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)); + + std::array local_correct{true, true}; + std::array elapsed_ms{}; + std::array actual{}; + for (int scenario = 0; scenario < 2; ++scenario) { + const bool in_place = scenario == 1; + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), send_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + float *active_recv = in_place ? d_send + local_offset : d_recv; + auto restore_local_block = [&]() { + if (in_place) { + CHECK_RT(Rt, Rt::Memcpy(active_recv, h_send.data() + local_offset, + recv_bytes, Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + } + }; + + for (int i = 0; i < kWarmupIterations; ++i) { + restore_local_block(); + CHECK_INFINI(infinicclReduceScatter(d_send, active_recv, kRecvCount, + infinicclFloat32, infinicclSum, comm, + nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + if (in_place) { + for (int i = 0; i < kProfileIterations; ++i) { + restore_local_block(); + Timer timer; + CHECK_INFINI(infinicclReduceScatter(d_send, active_recv, kRecvCount, + infinicclFloat32, infinicclSum, + comm, nullptr)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + elapsed_ms[scenario] += timer.ElapsedMs(); + } + elapsed_ms[scenario] /= static_cast(kProfileIterations); + } else { + Timer timer; + for (int i = 0; i < kProfileIterations; ++i) { + CHECK_INFINI(infinicclReduceScatter(d_send, active_recv, kRecvCount, + infinicclFloat32, infinicclSum, + comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + elapsed_ms[scenario] = + timer.ElapsedMs() / static_cast(kProfileIterations); + } + + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), active_recv, recv_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const float expected = ExpectedValue(rank, size); + local_correct[scenario] = + Validator::ValidateResult(h_recv.data(), kRecvCount, expected, rank); + actual[scenario] = h_recv.front(); + } + + // Reduce both validation flags into every destination block. Every rank then + // receives the same cluster-wide minimum and exits consistently. + if (world_size > std::numeric_limits::max() / 2 || + world_size * 2 > std::numeric_limits::max() / sizeof(int32_t)) { + std::cerr << "ReduceScatter status buffer size overflows `size_t`." + << std::endl; + std::exit(EXIT_FAILURE); + } + std::vector h_status_send(world_size * 2, 0); + for (int dest_rank = 0; dest_rank < size; ++dest_rank) { + const size_t offset = static_cast(dest_rank) * 2; + h_status_send[offset] = local_correct[0] ? 1 : 0; + h_status_send[offset + 1] = local_correct[1] ? 1 : 0; + } + std::array h_status_recv{}; + 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(infinicclReduceScatter(d_status_send, d_status_recv, + h_status_recv.size(), infinicclInt32, + infinicclMin, 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)); + + const std::array global_correct{h_status_recv[0] == 1, + h_status_recv[1] == 1}; + if (rank == 0) { + const float expected = ExpectedValue(rank, size); + PrintResult("Out-of-Place", global_correct[0], actual[0], expected, + elapsed_ms[0]); + PrintResult("In-Place", global_correct[1], actual[1], expected, + 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 ReduceScatter scenarios passed." + << std::endl; + } else { + std::cerr << "[Main Process] Hybrid ReduceScatter 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 RunReduceScatterExample(argc, argv) ? EXIT_SUCCESS : EXIT_FAILURE; +} diff --git a/examples/mpi/reduce_scatter.cc b/examples/mpi/reduce_scatter.cc index f1dbb4f..cb12fc9 100644 --- a/examples/mpi/reduce_scatter.cc +++ b/examples/mpi/reduce_scatter.cc @@ -6,7 +6,14 @@ #include +#include +#include +#include +#include +#include #include +#include +#include #include // Public API @@ -24,30 +31,57 @@ using namespace infini::ccl; -void RunReduceScatterExample(int argc, char **argv, int warmup_iter, - int profile_iter, const size_t kRecvCount) { +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 RunReduceScatterExample(int argc, char **argv, int warmup_iter, + int profile_iter, size_t recv_count) { 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 ReduceScatter." << std::endl; + std::exit(EXIT_FAILURE); + } - char hostname[256]; - gethostname(hostname, sizeof(hostname)); + 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)); - // 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); + std::array hostname{}; + if (gethostname(hostname.data(), hostname.size()) != 0) { + std::cerr << "Failed to query the hostname for ReduceScatter." << std::endl; + std::exit(EXIT_FAILURE); } + hostname.back() = '\0'; - std::cout << "[Rank " << rank << "] Host: " << hostname + std::cout << "[Rank " << rank << "] Host: " << hostname.data() << " | GPU: " << Device::StringFromType(kDevType) << " " << " | Device " << local_rank << std::endl; @@ -55,21 +89,34 @@ void RunReduceScatterExample(int argc, char **argv, int warmup_iter, infinicclComm_t comm = nullptr; CHECK_INFINI(infinicclCommInitAll(&comm, size, nullptr)); - // ReduceScatter requires `send_count = recv_count * world_size`. - const size_t kSendCount = kRecvCount * static_cast(size); + const size_t world_size = static_cast(size); + if (recv_count > std::numeric_limits::max() / world_size) { + std::cerr << "ReduceScatter element count overflows `size_t`." << std::endl; + std::exit(EXIT_FAILURE); + } + const size_t send_count = recv_count * world_size; + if (send_count > std::numeric_limits::max() / sizeof(float)) { + std::cerr << "ReduceScatter buffer size overflows `size_t`." << std::endl; + std::exit(EXIT_FAILURE); + } // Prepare Data - std::vector h_send(kSendCount); - std::vector h_recv(kRecvCount, 0.0f); - - // Initialize: each rank provides its (rank + 1) as data. - for (size_t i = 0; i < kSendCount; ++i) { - h_send[i] = static_cast(rank + 1); + std::vector h_send(send_count); + std::vector h_recv(recv_count, 0.0f); + + // Give each destination block a distinct reduced value so validation also + // checks that the scattered block matches this rank. + for (int dest_rank = 0; dest_rank < size; ++dest_rank) { + const float value = + static_cast(rank + 1) * static_cast(dest_rank + 1); + const size_t offset = static_cast(dest_rank) * recv_count; + std::fill_n(h_send.begin() + offset, recv_count, value); } - float *d_send, *d_recv; - size_t send_bytes = kSendCount * sizeof(*d_send); - size_t recv_bytes = kRecvCount * sizeof(*d_recv); + float *d_send = nullptr; + float *d_recv = nullptr; + size_t send_bytes = send_count * sizeof(*d_send); + size_t recv_bytes = recv_count * sizeof(*d_recv); CHECK_RT(Rt, Rt::Malloc((void **)&d_send, send_bytes)); CHECK_RT(Rt, Rt::Malloc((void **)&d_recv, recv_bytes)); @@ -81,9 +128,9 @@ void RunReduceScatterExample(int argc, char **argv, int warmup_iter, if (rank == 0) { std::cout << "\n=== Performing ReduceScatter on GPU Memory ===" << std::endl; - std::cout << "Recv data size per rank: " << kRecvCount << " floats (" + std::cout << "Recv data size per rank: " << recv_count << " floats (" << recv_bytes / 1024 / 1024 << " MB)" << std::endl; - std::cout << "Send data size per rank: " << kSendCount << " floats (" + std::cout << "Send data size per rank: " << send_count << " floats (" << send_bytes / 1024 / 1024 << " MB)" << std::endl; std::cout << "Operation: Sum" << std::endl; std::cout << "Warm-up iterations: " << warmup_iter << std::endl; @@ -93,14 +140,14 @@ void RunReduceScatterExample(int argc, char **argv, int warmup_iter, CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); // Warm-up and D2H transfer the answer. - CHECK_INFINI(infinicclReduceScatter(d_send, d_recv, kRecvCount, + CHECK_INFINI(infinicclReduceScatter(d_send, d_recv, recv_count, infinicclFloat32, infinicclSum, comm, nullptr)); CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, recv_bytes, Rt::MemcpyDeviceToHost)); for (int i = 1; i < warmup_iter; ++i) { - CHECK_INFINI(infinicclReduceScatter(d_send, d_recv, kRecvCount, + CHECK_INFINI(infinicclReduceScatter(d_send, d_recv, recv_count, infinicclFloat32, infinicclSum, comm, nullptr)); } @@ -110,7 +157,7 @@ void RunReduceScatterExample(int argc, char **argv, int warmup_iter, Timer timer; for (int i = 0; i < profile_iter; ++i) { - CHECK_INFINI(infinicclReduceScatter(d_send, d_recv, kRecvCount, + CHECK_INFINI(infinicclReduceScatter(d_send, d_recv, recv_count, infinicclFloat32, infinicclSum, comm, nullptr)); } @@ -122,13 +169,12 @@ void RunReduceScatterExample(int argc, char **argv, int warmup_iter, double elapsed = timer.ElapsedMs() / static_cast(profile_iter); // Result Validation: - float expected = 0.0f; - for (int r = 0; r < size; ++r) { - expected += static_cast(r + 1); - } + const float rank_sum = + static_cast(size) * (static_cast(size) + 1.0f) / 2.0f; + const float expected = rank_sum * static_cast(rank + 1); - Validator::ValidateResult(h_recv.data(), kRecvCount, expected, rank, true, - "ReduceScatter"); + const bool correct = Validator::ValidateResult( + h_recv.data(), recv_count, expected, rank, true, "ReduceScatter"); // Metrics Reporting (Only from rank 0 for cleaner output) if (rank == 0) { @@ -146,14 +192,17 @@ void RunReduceScatterExample(int argc, char **argv, int warmup_iter, if (rank == 0) { std::cout << "InfiniCCL finalized." << std::endl; } + return correct; } +} // namespace + int main(int argc, char **argv) { int warmup_iters = 2; int profile_iters = 20; size_t recv_count = 1 << 20; - RunReduceScatterExample(argc, argv, warmup_iters, profile_iters, recv_count); - - return EXIT_SUCCESS; + const bool correct = RunReduceScatterExample(argc, argv, warmup_iters, + profile_iters, recv_count); + return correct ? EXIT_SUCCESS : EXIT_FAILURE; } diff --git a/src/backends/ccl/common/impl/reduce_scatter.h b/src/backends/ccl/common/impl/reduce_scatter.h new file mode 100644 index 0000000..ba4b617 --- /dev/null +++ b/src/backends/ccl/common/impl/reduce_scatter.h @@ -0,0 +1,69 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_REDUCE_SCATTER_H_ +#define INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_REDUCE_SCATTER_H_ + +#include "backends/ccl/common/api.h" +#include "backends/ccl/common/comm_instance.h" +#include "base/reduce_scatter.h" +#include "communicator.h" + +namespace infini::ccl { + +template +struct DeferredReduceScatter { + using type = ReduceScatter; +}; + +template +class CclReduceScatterImpl { + public: + static ReturnStatus Apply(const void *send_buff, void *recv_buff, + size_t recv_count, DataType data_type, + ReductionOpType op, 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 DeferredReduceScatter::type; + if constexpr (BackendEnabled::value) { + return ReduceScatterImpl::Apply( + send_buff, recv_buff, recv_count, data_type, op, 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; + } + + typename Api::RedOp native_op{}; + if (!TypeMap::ToBackendRedOp(op, &native_op)) { + return ReturnStatus::kNotSupported; + } + + return Api::Check(Api::ReduceScatter( + send_buff, recv_buff, recv_count, native_type, native_op, + instance->handle, reinterpret_cast(stream))); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_REDUCE_SCATTER_H_ diff --git a/src/backends/ccl/mccl/api.h b/src/backends/ccl/mccl/api.h index a5d2bcd..e864771 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 ReduceScatter(const void *send_buff, void *recv_buff, + size_t recv_count, DataType data_type, RedOp op, + Comm comm, Stream stream) { + return mcclReduceScatter(send_buff, recv_buff, recv_count, data_type, op, + comm, stream); + } }; } // namespace infini::ccl diff --git a/src/backends/ccl/mccl/impl/reduce_scatter.h b/src/backends/ccl/mccl/impl/reduce_scatter.h new file mode 100644 index 0000000..6e9614e --- /dev/null +++ b/src/backends/ccl/mccl/impl/reduce_scatter.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_REDUCE_SCATTER_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_REDUCE_SCATTER_H_ + +#include "backends/ccl/common/impl/reduce_scatter.h" + +namespace infini::ccl { + +template +class ReduceScatterImpl + : public CclReduceScatterImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_REDUCE_SCATTER_H_ diff --git a/src/backends/ccl/nccl/api.h b/src/backends/ccl/nccl/api.h index e7b6119..21e5825 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 ReduceScatter(const void *send_buff, void *recv_buff, + size_t recv_count, DataType data_type, RedOp op, + Comm comm, Stream stream) { + return ncclReduceScatter(send_buff, recv_buff, recv_count, data_type, op, + comm, stream); + } }; } // namespace infini::ccl diff --git a/src/backends/ccl/nccl/impl/reduce_scatter.h b/src/backends/ccl/nccl/impl/reduce_scatter.h new file mode 100644 index 0000000..f334198 --- /dev/null +++ b/src/backends/ccl/nccl/impl/reduce_scatter.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_REDUCE_SCATTER_H_ +#define INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_REDUCE_SCATTER_H_ + +#include "backends/ccl/common/impl/reduce_scatter.h" + +namespace infini::ccl { + +template +class ReduceScatterImpl + : public CclReduceScatterImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_REDUCE_SCATTER_H_ diff --git a/src/backends/mpi/ompi/impl/reduce_scatter.h b/src/backends/mpi/ompi/impl/reduce_scatter.h index 5c8c0a9..98bbadb 100644 --- a/src/backends/mpi/ompi/impl/reduce_scatter.h +++ b/src/backends/mpi/ompi/impl/reduce_scatter.h @@ -1,7 +1,9 @@ #ifndef INFINI_CCL_BACKENDS_MPI_OMPI_IMPL_REDUCE_SCATTER_H_ #define INFINI_CCL_BACKENDS_MPI_OMPI_IMPL_REDUCE_SCATTER_H_ +#include #include +#include #include #include "backends/mpi/ompi/checks.h" @@ -31,42 +33,58 @@ class ReduceScatterImpl { LOG("Invalid OpenMPI communicator instance for `ReduceScatter`."); return ReturnStatus::kInternalError; } + if (comm->size() <= 0) { + LOG("Invalid world size for `ReduceScatter`."); + return ReturnStatus::kInternalError; + } MPI_Datatype mpi_type = DataTypeToOmpiType(data_type); MPI_Op mpi_op = RedOpToOmpiOp(op); + if (mpi_type == MPI_BYTE) { + LOG("Data type is not supported by OpenMPI reductions for " + "`ReduceScatter`."); + return ReturnStatus::kNotSupported; + } size_t world_size = static_cast(comm->size()); size_t type_size = kDataTypeToSize.at(data_type); + if (recv_count > std::numeric_limits::max() / world_size) { + LOG("Send element count overflows `size_t` for `ReduceScatter`."); + return ReturnStatus::kInvalidArgument; + } size_t send_count = recv_count * world_size; + if (send_count > std::numeric_limits::max() / type_size || + recv_count > std::numeric_limits::max() / type_size) { + LOG("Buffer byte size overflows `size_t` for `ReduceScatter`."); + return ReturnStatus::kInvalidArgument; + } size_t send_bytes = send_count * type_size; size_t recv_bytes = recv_count * type_size; + if (recv_count > static_cast(std::numeric_limits::max())) { + LOG("Receive count exceeds MPI int range for `ReduceScatter`."); + return ReturnStatus::kInvalidArgument; + } + int mpi_recv_count = static_cast(recv_count); + // 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); + std::unique_ptr host_sendbuf( + std::malloc(send_bytes), &std::free); + std::unique_ptr host_recvbuf( + std::malloc(recv_bytes), &std::free); if (!host_sendbuf || !host_recvbuf) { - free(host_sendbuf); - free(host_recvbuf); LOG("Failed to allocate host buffers for `ReduceScatter` 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))); - if (recv_count > static_cast(std::numeric_limits::max())) { - LOG("recv_count exceeds MPI int range for `ReduceScatter`."); - free(host_sendbuf); - free(host_recvbuf); - return ReturnStatus::kInvalidArgument; - } - int mpi_recv_count = static_cast(recv_count); - - INFINI_CHECK_MPI(MPI_Reduce_scatter_block(host_sendbuf, host_recvbuf, - mpi_recv_count, mpi_type, mpi_op, - inst->handle)); + INFINI_CHECK_MPI(MPI_Reduce_scatter_block( + host_sendbuf.get(), host_recvbuf.get(), mpi_recv_count, mpi_type, + mpi_op, inst->handle)); if (op == ReductionOpType::kAvg) { float scale = 1.0f / static_cast(world_size); @@ -74,7 +92,7 @@ class ReduceScatterImpl { DispatchFunc(data_type, [&](auto dtype) { using T = typename decltype(dtype)::type; - T *typed_buf = static_cast(host_recvbuf); + T *typed_buf = static_cast(host_recvbuf.get()); // Simply do the averaging on the CPU before the H2D copy. for (size_t i = 0; i < recv_count; ++i) { @@ -94,12 +112,9 @@ class ReduceScatterImpl { }); } - 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)); - free(host_sendbuf); - free(host_recvbuf); - return ReturnStatus::kSuccess; } }; diff --git a/src/base/reduce_scatter.h b/src/base/reduce_scatter.h index c06150b..6b1322c 100644 --- a/src/base/reduce_scatter.h +++ b/src/base/reduce_scatter.h @@ -20,9 +20,14 @@ class ReduceScatter : public Operation { size_t recv_count, DataType datatype, ReductionOpType op, void *comm_handle, void *stream) { - if (HasInvalidArgs(send_buff, recv_buff, datatype, op, comm_handle)) { + if (HasInvalidArgs(send_buff, recv_buff, recv_count, datatype, op, + comm_handle)) { return ReturnStatus::kInvalidArgument; } + if (recv_count == 0) { + return ReturnStatus::kSuccess; + } + auto *comm = static_cast(comm_handle); return ReduceScatterImpl::Apply( send_buff, recv_buff, recv_count, datatype, op, comm, stream); @@ -30,17 +35,13 @@ class ReduceScatter : public Operation { private: static bool HasInvalidArgs(const void *send_buff, void *recv_buff, - DataType datatype, ReductionOpType op, - void *comm_handle) { + size_t recv_count, DataType datatype, + ReductionOpType op, void *comm_handle) { if (!comm_handle) { // TODO(lzm): change to use `glog`. LOG("Invalid communicator handle for `ReduceScatter`."); return true; } - if (!send_buff || !recv_buff) { - LOG("Invalid buffer pointer for `ReduceScatter`."); - return true; - } if (op < ReductionOpType::kSum || op >= ReductionOpType::kNumRedOps) { LOG("Invalid reduction operation for `ReduceScatter`."); return true; @@ -49,6 +50,13 @@ class ReduceScatter : public Operation { LOG("Invalid data type for `ReduceScatter`."); return true; } + if (recv_count == 0) { + return false; + } + if (!send_buff || !recv_buff) { + LOG("Invalid buffer pointer for `ReduceScatter`."); + return true; + } return false; } };