diff --git a/examples/ccl/reduce.cc b/examples/ccl/reduce.cc new file mode 100644 index 0000000..c857ebd --- /dev/null +++ b/examples/ccl/reduce.cc @@ -0,0 +1,307 @@ +/** + * InfiniCCL Example: Thread-per-GPU Single-Node Reduce + * + * This example creates one native CCL rank per GPU and validates a rooted + * reduction without an MPI launcher. + */ + +#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 RunState { + std::atomic correct{true}; + std::atomic completed{0}; +}; + +struct ThreadArgs { + int rank; + int size; + infinicclUniqueId id; + size_t count; + int warmup_iterations; + int profile_iterations; + RunState *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; +} + +float ExpectedValue(int world_size) { + return static_cast(world_size) * + (static_cast(world_size) + 1.0f) / 2.0f; +} + +void PrintReduceMetrics(size_t count, double elapsed_ms) { + constexpr double kBytesPerMiB = 1024.0 * 1024.0; + constexpr double kBytesPerGB = 1.0e9; + const double payload_bytes = static_cast(count) * sizeof(float); + const auto original_flags = std::cout.flags(); + const auto original_precision = std::cout.precision(); + + std::cout << "Data size: " << count << " floats (" << std::fixed + << std::setprecision(2) << payload_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 = + payload_bytes / kBytesPerGB / (elapsed_ms / 1000.0); + // `nccl-tests` uses a bus-bandwidth correction factor of 1 for Reduce. + const double bus_bandwidth = algorithm_bandwidth; + std::cout << "Throughput: " << std::setprecision(2) << bus_bandwidth + << " GB/s (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: " << std::setprecision(2) + << 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, float expected, float actual, size_t count, + 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 Reduce 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 << "Expect: " << expected << std::endl; + std::cout << "Actual: " << actual << std::endl; + PrintReduceMetrics(count, elapsed_ms); +} + +void WaitForAll(RunState *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 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 Reduce 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 bool is_root = args.rank == kRoot; + const size_t total_bytes = args.count * sizeof(float); + std::vector h_send(args.count, static_cast(args.rank + 1)); + std::vector h_recv(is_root ? args.count : 0, 0.0f); + + float *d_send = nullptr; + float *d_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), total_bytes)); + if (is_root) { + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_recv), total_bytes)); + } + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + for (int i = 0; i < args.warmup_iterations; ++i) { + CHECK_INFINI(infinicclReduce(d_send, d_recv, args.count, infinicclFloat32, + infinicclSum, kRoot, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + Timer timer; + for (int i = 0; i < args.profile_iterations; ++i) { + CHECK_INFINI(infinicclReduce(d_send, d_recv, args.count, infinicclFloat32, + infinicclSum, kRoot, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const double elapsed_ms = + timer.ElapsedMs() / static_cast(args.profile_iterations); + + bool local_correct = true; + const float expected = ExpectedValue(args.size); + if (is_root) { + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + local_correct = Validator::ValidateResult( + h_recv.data(), args.count, expected, args.rank, false, "Reduce"); + if (!local_correct) { + args.state->correct.store(false, std::memory_order_relaxed); + } + } + + WaitForAll(args.state, args.size); + if (is_root) { + PrintResult(args.state->correct.load(std::memory_order_acquire), expected, + h_recv.front(), args.count, elapsed_ms); + } + + CHECK_RT(Rt, Rt::Free(d_send)); + if (is_root) { + 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 Warm-up iterations (default: 2)\n" + << " -p Profile iterations (default: 20)\n" + << " -n Elements reduced per 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 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_iterations); + break; + case 'p': + parsed = ParsePositiveNumber(optarg, &profile_iterations); + break; + case 'n': + parsed = ParsePositiveNumber(optarg, &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 Reduce." << std::endl; + return EXIT_FAILURE; + } + } + + if (optind != argc) { + std::cerr << "Unexpected positional argument for Reduce." << std::endl; + return EXIT_FAILURE; + } + if (count > std::numeric_limits::max() / sizeof(float)) { + std::cerr << "Reduce 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 Reduce." << 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] Count: " << count + << " floats | Warmup: " << warmup_iterations + << " | Profile: " << profile_iterations << " | Root: " << kRoot + << std::endl; + + infinicclUniqueId shared_id{}; + CHECK_INFINI(infinicclGetUniqueId(&shared_id)); + + RunState state; + std::vector threads; + threads.reserve(num_gpus); + for (int rank = 0; rank < num_gpus; ++rank) { + ThreadArgs args{rank, num_gpus, shared_id, + count, 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 Reduce validation passed." << std::endl; + } else { + std::cerr << "[Main Process] CCL Reduce 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.cc b/examples/ccl_mpi_hybrid/reduce.cc new file mode 100644 index 0000000..f043ceb --- /dev/null +++ b/examples/ccl_mpi_hybrid/reduce.cc @@ -0,0 +1,276 @@ +/** + * InfiniCCL Example: Reduce (OpenMPI + CCL Hybrid) + * + * This example uses an OpenMPI inter communicator to distribute a native CCL + * unique ID, then validates a rooted GPU reduction on the native communicator. + */ + +#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; +} + +float ExpectedValue(int world_size) { + return static_cast(world_size) * + (static_cast(world_size) + 1.0f) / 2.0f; +} + +void PrintReduceMetrics(size_t count, double elapsed_ms) { + constexpr double kBytesPerMiB = 1024.0 * 1024.0; + constexpr double kBytesPerGB = 1.0e9; + const double payload_bytes = static_cast(count) * sizeof(float); + const auto original_flags = std::cout.flags(); + const auto original_precision = std::cout.precision(); + + std::cout << "Data size: " << count << " floats (" << std::fixed + << std::setprecision(2) << payload_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 = + payload_bytes / kBytesPerGB / (elapsed_ms / 1000.0); + // `nccl-tests` uses a bus-bandwidth correction factor of 1 for Reduce. + const double bus_bandwidth = algorithm_bandwidth; + std::cout << "Throughput: " << std::setprecision(2) << bus_bandwidth + << " GB/s (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: " << std::setprecision(2) + << 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, float expected, float actual, size_t count, + 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 Reduce 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 << "Expect: " << expected << std::endl; + std::cout << "Actual: " << actual << std::endl; + PrintReduceMetrics(count, elapsed_ms); +} + +bool RunReduceExample(int argc, char **argv) { + constexpr Device::Type kDevType = + ListGetBest(EnabledDevices{}); + using Rt = Runtime; + + constexpr int kWarmupIterations = 2; + constexpr int kProfileIterations = 20; + constexpr size_t kCount = 1 << 20; + + CHECK_INFINI(infinicclInit(&argc, &argv)); + + int rank = -1; + int size = 0; + CHECK_INFINI(infinicclGetRank(&rank)); + CHECK_INFINI(infinicclGetSize(&size)); + if (size <= 0 || kRoot >= size) { + std::cerr << "Invalid world size for hybrid Reduce." << 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 Reduce." << 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)); + + // Exercise Reduce before the native communicator exists. The CCL backend + // must delegate this call to the OpenMPI inter communicator. + constexpr size_t kBootstrapCount = 1; + const float bootstrap_send = static_cast(rank + 1); + float bootstrap_recv = 0.0f; + float *d_bootstrap_send = nullptr; + float *d_bootstrap_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_bootstrap_send), + kBootstrapCount * sizeof(float))); + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_bootstrap_recv), + kBootstrapCount * sizeof(float))); + } + CHECK_RT(Rt, Rt::Memcpy(d_bootstrap_send, &bootstrap_send, sizeof(float), + Rt::MemcpyHostToDevice)); + CHECK_INFINI(infinicclReduce(d_bootstrap_send, d_bootstrap_recv, + kBootstrapCount, infinicclFloat32, infinicclSum, + kRoot, comm, nullptr)); + if (rank == kRoot) { + CHECK_RT(Rt, Rt::Memcpy(&bootstrap_recv, d_bootstrap_recv, sizeof(float), + Rt::MemcpyDeviceToHost)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + int32_t bootstrap_status = + rank != kRoot || bootstrap_recv == ExpectedValue(size) ? 1 : 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)); + } + + infinicclUniqueId id{}; + if (rank == kRoot) { + CHECK_INFINI(infinicclGetUniqueId(&id)); + } + CHECK_INFINI(infinicclBroadcast(&id, &id, sizeof(id), infinicclChar, kRoot, + comm, nullptr)); + CHECK_INFINI(infinicclCommInitRank(&comm, size, id, rank)); + + if (kCount > std::numeric_limits::max() / sizeof(float)) { + std::cerr << "Hybrid Reduce buffer size overflows `size_t`." << std::endl; + std::exit(EXIT_FAILURE); + } + const bool is_root = rank == kRoot; + const size_t total_bytes = kCount * sizeof(float); + std::vector h_send(kCount, static_cast(rank + 1)); + std::vector h_recv(is_root ? kCount : 0, 0.0f); + + float *d_send = nullptr; + float *d_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), total_bytes)); + if (is_root) { + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_recv), total_bytes)); + } + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + if (is_root) { + std::cout << "\n=== Performing Hybrid CCL Reduce on GPU Memory ===" + << std::endl; + std::cout << "Operation: Sum" << std::endl; + std::cout << "Root rank: " << kRoot << std::endl; + std::cout << "Warm-up iterations: " << kWarmupIterations << std::endl; + std::cout << "Profile iterations: " << kProfileIterations << std::endl; + } + + for (int i = 0; i < kWarmupIterations; ++i) { + CHECK_INFINI(infinicclReduce(d_send, d_recv, kCount, infinicclFloat32, + infinicclSum, kRoot, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + Timer timer; + for (int i = 0; i < kProfileIterations; ++i) { + CHECK_INFINI(infinicclReduce(d_send, d_recv, kCount, infinicclFloat32, + infinicclSum, kRoot, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const double elapsed_ms = + timer.ElapsedMs() / static_cast(kProfileIterations); + + bool correct = bootstrap_status == 1; + const float expected = ExpectedValue(size); + if (is_root) { + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + correct = Validator::ValidateResult(h_recv.data(), kCount, expected, rank, + false, "Reduce") && + correct; + PrintResult(correct, expected, h_recv.front(), kCount, elapsed_ms); + } + + int32_t validation_status = correct ? 1 : 0; + CHECK_INFINI(infinicclBroadcast(&validation_status, &validation_status, 1, + infinicclInt32, kRoot, comm, nullptr)); + correct = validation_status == 1; + + CHECK_RT(Rt, Rt::Free(d_send)); + if (is_root) { + CHECK_RT(Rt, Rt::Free(d_recv)); + } + CHECK_INFINI(infinicclCommDestroy(comm)); + CHECK_INFINI(infinicclFinalize()); + + if (is_root) { + if (correct) { + std::cout << "[Main Process] Hybrid CCL Reduce validation passed." + << std::endl; + } else { + std::cerr << "[Main Process] Hybrid CCL Reduce validation failed." + << std::endl; + } + std::cout << "InfiniCCL finalized." << std::endl; + } + return correct; +} + +} // namespace + +int main(int argc, char **argv) { + return RunReduceExample(argc, argv) ? EXIT_SUCCESS : EXIT_FAILURE; +} diff --git a/examples/mpi/reduce.cc b/examples/mpi/reduce.cc index a998b7f..f361d92 100644 --- a/examples/mpi/reduce.cc +++ b/examples/mpi/reduce.cc @@ -1,14 +1,24 @@ /** - * InfiniCCL Example: Reduce + * InfiniCCL Example: Reduce (MPI Backend) * - * This example demonstrates the API for performing a collective - * reduction across multiple accelerators and nodes, where only - * `root` receives the reduced result. + * This example performs a rooted reduction across multiple accelerators and + * nodes, validates the root result, and propagates validation failure through + * the process exit status. */ #include +#include +#include +#include +#include +#include +#include +#include #include +#include +#include +#include #include // Public API @@ -26,127 +36,196 @@ using namespace infini::ccl; -void RunReduceExample(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; +} + +float ExpectedValue(int world_size) { + return static_cast(world_size) * + (static_cast(world_size) + 1.0f) / 2.0f; +} + +void PrintReduceMetrics(size_t count, double elapsed_ms) { + constexpr double kBytesPerMiB = 1024.0 * 1024.0; + constexpr double kBytesPerGB = 1.0e9; + const double payload_bytes = static_cast(count) * sizeof(float); + const auto original_flags = std::cout.flags(); + const auto original_precision = std::cout.precision(); + + std::cout << "Data size: " << count << " floats (" << std::fixed + << std::setprecision(2) << payload_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 = + payload_bytes / kBytesPerGB / (elapsed_ms / 1000.0); + // `nccl-tests` uses a bus-bandwidth correction factor of 1 for Reduce. + const double bus_bandwidth = algorithm_bandwidth; + std::cout << "Throughput: " << std::setprecision(2) << bus_bandwidth + << " GB/s (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: " << std::setprecision(2) + << 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, float expected, float actual, size_t count, + double elapsed_ms) { + constexpr const char *kGreen = "\033[32m"; + constexpr const char *kRed = "\033[31m"; + constexpr const char *kReset = "\033[0m"; + + std::cout << "\n=== MPI Reduce 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 << "Expect: " << expected << std::endl; + std::cout << "Actual: " << actual << std::endl; + PrintReduceMetrics(count, elapsed_ms); +} + +bool RunReduceExample(int argc, char **argv, int warmup_iterations, + int profile_iterations, size_t 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 || kRoot >= size) { + std::cerr << "Invalid world size for MPI Reduce." << 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 Reduce." << 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 Reduce - constexpr int kRoot = 0; - - // Prepare Data - std::vector h_send(kNumElements); - std::vector h_recv(kNumElements, 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); + if (count == 0 || + count > std::numeric_limits::max() / sizeof(float)) { + std::cerr << "Invalid MPI Reduce buffer size." << std::endl; + std::exit(EXIT_FAILURE); + } + const bool is_root = rank == kRoot; + const size_t total_bytes = count * sizeof(float); + std::vector h_send(count, static_cast(rank + 1)); + std::vector h_recv(is_root ? count : 0, 0.0f); + + float *d_send = nullptr; + float *d_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), total_bytes)); + if (is_root) { + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_recv), total_bytes)); } - - float *d_send, *d_recv; - size_t total_bytes = kNumElements * sizeof(*d_send); - CHECK_RT(Rt, Rt::Malloc((void **)&d_send, total_bytes)); - CHECK_RT(Rt, Rt::Malloc((void **)&d_recv, total_bytes)); CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes, Rt::MemcpyHostToDevice)); - CHECK_RT(Rt, Rt::Memcpy(d_recv, h_recv.data(), total_bytes, - Rt::MemcpyHostToDevice)); - - if (rank == kRoot) { - std::cout << "\n=== Performing Reduce on GPU Memory ===" << std::endl; - std::cout << "Data size: " << kNumElements << " floats (" - << total_bytes / 1024 / 1024 << " MB)" << std::endl; - std::cout << "Operation: Sum" << 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::StreamSynchronize(nullptr)); - // Warm-up and D2H transfer the answer on `root`. - CHECK_INFINI(infinicclReduce(d_send, d_recv, kNumElements, infinicclFloat32, - infinicclSum, kRoot, comm, nullptr)); - if (rank == kRoot) { - CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes, - Rt::MemcpyDeviceToHost)); + if (is_root) { + std::cout << "\n=== Performing MPI Reduce on GPU Memory ===" << std::endl; + std::cout << "Operation: Sum" << std::endl; + std::cout << "Root rank: " << kRoot << 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(infinicclReduce(d_send, d_recv, kNumElements, infinicclFloat32, + for (int i = 0; i < warmup_iterations; ++i) { + CHECK_INFINI(infinicclReduce(d_send, d_recv, count, infinicclFloat32, infinicclSum, kRoot, comm, nullptr)); } CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); - // Profiling Timer timer; - - for (int i = 0; i < profile_iter; ++i) { - CHECK_INFINI(infinicclReduce(d_send, d_recv, kNumElements, infinicclFloat32, + for (int i = 0; i < profile_iterations; ++i) { + CHECK_INFINI(infinicclReduce(d_send, d_recv, count, infinicclFloat32, infinicclSum, 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`). - if (rank == kRoot) { - float expected = 0.0f; - for (int r = 0; r < size; ++r) { - expected += static_cast(r + 1); - } - - Validator::ValidateResult(h_recv.data(), kNumElements, expected, rank, true, - "Reduce"); - - Metrics metrics{elapsed, total_bytes, size}; - metrics.Print(); + bool correct = true; + const float expected = ExpectedValue(size); + if (is_root) { + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + correct = Validator::ValidateResult(h_recv.data(), count, expected, rank, + false, "Reduce"); + PrintResult(correct, expected, h_recv.front(), count, elapsed_ms); } - // Cleanup - CHECK_RT(Rt, Rt::Free(d_send)); - CHECK_RT(Rt, Rt::Free(d_recv)); + int32_t validation_status = correct ? 1 : 0; + CHECK_INFINI(infinicclBroadcast(&validation_status, &validation_status, 1, + infinicclInt32, kRoot, comm, nullptr)); + correct = validation_status == 1; + CHECK_RT(Rt, Rt::Free(d_send)); + if (is_root) { + CHECK_RT(Rt, Rt::Free(d_recv)); + } CHECK_INFINI(infinicclCommDestroy(comm)); CHECK_INFINI(infinicclFinalize()); - if (rank == kRoot) { + if (is_root) { 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; +} // namespace - RunReduceExample(argc, argv, warmup_iters, profile_iters, num_elements); - - return EXIT_SUCCESS; +int main(int argc, char **argv) { + constexpr int kWarmupIterations = 2; + constexpr int kProfileIterations = 20; + constexpr size_t kCount = 1 << 20; + return RunReduceExample(argc, argv, kWarmupIterations, kProfileIterations, + kCount) + ? EXIT_SUCCESS + : EXIT_FAILURE; } diff --git a/src/backends/ccl/common/impl/reduce.h b/src/backends/ccl/common/impl/reduce.h new file mode 100644 index 0000000..a20b20a --- /dev/null +++ b/src/backends/ccl/common/impl/reduce.h @@ -0,0 +1,76 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_REDUCE_H_ +#define INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_REDUCE_H_ + +#include "backends/ccl/common/api.h" +#include "backends/ccl/common/comm_instance.h" +#include "base/reduce.h" +#include "communicator.h" +#include "logging.h" + +namespace infini::ccl { + +template +struct DeferredReduce { + using type = Reduce; +}; + +template +class CclReduceImpl { + public: + static ReturnStatus Apply(const void *send_buff, void *recv_buff, + size_t count, DataType data_type, + ReductionOpType op, 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 DeferredReduce::type; + if constexpr (BackendEnabled::value) { + return ReduceImpl::Apply( + send_buff, recv_buff, count, data_type, op, 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 `Reduce`."); + 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::Reduce( + send_buff, recv_buff, count, native_type, native_op, root, + instance->handle, reinterpret_cast(stream))); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_REDUCE_H_ diff --git a/src/backends/ccl/mccl/api.h b/src/backends/ccl/mccl/api.h index a5d2bcd..97c3a08 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 Reduce(const void *send_buff, void *recv_buff, size_t count, + DataType data_type, RedOp op, int root, Comm comm, + Stream stream) { + return mcclReduce(send_buff, recv_buff, count, data_type, op, root, comm, + stream); + } }; } // namespace infini::ccl diff --git a/src/backends/ccl/mccl/impl/reduce.h b/src/backends/ccl/mccl/impl/reduce.h new file mode 100644 index 0000000..e8cf6fb --- /dev/null +++ b/src/backends/ccl/mccl/impl/reduce.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_REDUCE_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_REDUCE_H_ + +#include "backends/ccl/common/impl/reduce.h" + +namespace infini::ccl { + +template +class ReduceImpl + : public CclReduceImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_REDUCE_H_ diff --git a/src/backends/ccl/nccl/api.h b/src/backends/ccl/nccl/api.h index e7b6119..c547808 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 Reduce(const void *send_buff, void *recv_buff, size_t count, + DataType data_type, RedOp op, int root, Comm comm, + Stream stream) { + return ncclReduce(send_buff, recv_buff, count, data_type, op, root, comm, + stream); + } }; } // namespace infini::ccl diff --git a/src/backends/ccl/nccl/impl/reduce.h b/src/backends/ccl/nccl/impl/reduce.h new file mode 100644 index 0000000..3d22648 --- /dev/null +++ b/src/backends/ccl/nccl/impl/reduce.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_REDUCE_H_ +#define INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_REDUCE_H_ + +#include "backends/ccl/common/impl/reduce.h" + +namespace infini::ccl { + +template +class ReduceImpl + : public CclReduceImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_REDUCE_H_ diff --git a/src/backends/mpi/ompi/impl/reduce.h b/src/backends/mpi/ompi/impl/reduce.h index c50bfda..d58671a 100644 --- a/src/backends/mpi/ompi/impl/reduce.h +++ b/src/backends/mpi/ompi/impl/reduce.h @@ -3,6 +3,7 @@ #include #include +#include #include #include "backends/mpi/ompi/checks.h" @@ -28,14 +29,27 @@ class ReduceImpl { ListGetBest(ActiveDevices{}); using Rt = Runtime; + if (!comm || comm->inter_comm_backend() != BackendType::kOmpi) { + LOG("Invalid OpenMPI communicator for `Reduce`."); + return ReturnStatus::kInternalError; + } + auto *inst = static_cast(comm->inter_comm()); if (!inst || inst->handle == MPI_COMM_NULL) { LOG("Invalid OpenMPI communicator instance for `Reduce`."); return ReturnStatus::kInternalError; } + if (comm->size() <= 0) { + LOG("Invalid world size for `Reduce`."); + 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 `Reduce`."); + return ReturnStatus::kNotSupported; + } if (count > static_cast(std::numeric_limits::max())) { LOG("`count` exceeds MPI int range for `Reduce`."); @@ -44,26 +58,31 @@ class ReduceImpl { int mpi_count = static_cast(count); size_t type_size = kDataTypeToSize.at(data_type); + if (count > std::numeric_limits::max() / type_size) { + LOG("Buffer byte size overflows `size_t` for `Reduce`."); + return ReturnStatus::kInvalidArgument; + } size_t total_bytes = count * type_size; const bool is_root = comm->rank() == root; // Host staging buffers. Only `root` allocates the receive side, since // `MPI_Reduce` writes the output only on `root`. - void *host_sendbuf = std::malloc(total_bytes); - void *host_recvbuf = is_root ? std::malloc(total_bytes) : nullptr; + std::unique_ptr host_sendbuf( + std::malloc(total_bytes), &std::free); + std::unique_ptr host_recvbuf( + is_root ? std::malloc(total_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 `Reduce` staging."); return ReturnStatus::kSystemError; } - CHECK_STATUS(Rt, Rt::Memcpy(host_sendbuf, send_buff, total_bytes, + CHECK_STATUS(Rt, Rt::Memcpy(host_sendbuf.get(), send_buff, total_bytes, Rt::MemcpyDeviceToHost)); CHECK_STATUS(Rt, Rt::StreamSynchronize(static_cast(stream))); - INFINI_CHECK_MPI(MPI_Reduce(host_sendbuf, host_recvbuf, mpi_count, mpi_type, - mpi_op, root, inst->handle)); + INFINI_CHECK_MPI(MPI_Reduce(host_sendbuf.get(), host_recvbuf.get(), + mpi_count, mpi_type, mpi_op, root, + inst->handle)); if (is_root) { if (op == ReductionOpType::kAvg) { @@ -72,7 +91,7 @@ class ReduceImpl { 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 < count; ++i) { @@ -92,13 +111,10 @@ class ReduceImpl { }); } - CHECK_STATUS(Rt, Rt::Memcpy(recv_buff, host_recvbuf, total_bytes, + CHECK_STATUS(Rt, Rt::Memcpy(recv_buff, host_recvbuf.get(), total_bytes, Rt::MemcpyHostToDevice)); } - std::free(host_sendbuf); - std::free(host_recvbuf); - return ReturnStatus::kSuccess; } }; diff --git a/src/base/reduce.h b/src/base/reduce.h index 4415c2c..5cd5cbf 100644 --- a/src/base/reduce.h +++ b/src/base/reduce.h @@ -26,7 +26,7 @@ class Reduce : public Operation { } auto *comm = static_cast(comm_handle); - if (HasInvalidArgs(send_buff, recv_buff, datatype, op, root, comm)) { + if (HasInvalidArgs(send_buff, recv_buff, count, datatype, op, root, comm)) { return ReturnStatus::kInvalidArgument; } if (count == 0) { @@ -39,8 +39,8 @@ class Reduce : public Operation { private: static bool HasInvalidArgs(const void *send_buff, void *recv_buff, - DataType datatype, ReductionOpType op, int root, - Communicator *comm) { + size_t count, DataType datatype, + ReductionOpType op, int root, Communicator *comm) { if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) { LOG("Invalid data type for `Reduce`."); return true; @@ -53,6 +53,9 @@ class Reduce : public Operation { LOG("Invalid root rank for `Reduce`."); return true; } + if (count == 0) { + return false; + } if (!send_buff) { LOG("Invalid send buffer pointer for `Reduce`."); return true;