diff --git a/examples/ccl/reduce_scatter.cc b/examples/ccl/reduce_scatter.cc new file mode 100644 index 0000000..55a0210 --- /dev/null +++ b/examples/ccl/reduce_scatter.cc @@ -0,0 +1,257 @@ +/** + * InfiniCCL Example: Thread-per-GPU Single-Node `ReduceScatter` + * + * This example spawns one CPU thread per GPU on a single node, generates a + * shared UniqueID, and performs `ReduceScatter` entirely through InfiniCCL's + * native CCL backend without any OpenMPI dependency. + * + * Every rank's send buffer is split into `size` equally sized blocks; block `b` + * of rank `r` holds `(r + 1) + b`. After the reduction, rank `r` receives block + * `r` of the element-wise sum, so each block value is validated independently + * to catch misplaced destination chunks. + * + * Note: to properly run this example, you should specify `--launcher none` + * since it requires no MPI process spawning. + */ + +#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 { + +// Structure to pass execution data to each GPU worker thread. +struct ThreadArgs { + int rank; + int size; + infinicclUniqueId id; + size_t recv_count; // Elements received per rank. + int warmup_iter; + int profile_iter; + std::atomic_bool* all_correct; +}; + +// Value stored in block `block` of rank `rank`'s send buffer. The block index +// is embedded in the payload so a block mixed up by the backend cannot pass +// validation. +float SeedValue(int rank, size_t block) { + return static_cast(rank + 1) + static_cast(block); +} + +// Value block `block` holds after the element-wise sum over all ranks. +float ExpectedValue(int size, size_t block) { + const float rank_sum = static_cast(size * (size + 1) / 2); + return rank_sum + static_cast(size) * static_cast(block); +} + +// Every rank checks the single reduced block it received from the collective. +bool Validate(const std::vector& output, int size, int rank) { + return Validator::ValidateResult( + output.data(), output.size(), + ExpectedValue(size, static_cast(rank)), rank); +} + +void PrintResult(const std::vector& output, int rank, int size, + bool correct) { + const char* green = "\033[32m"; + const char* red = "\033[31m"; + const char* reset = "\033[0m"; + + std::cout << "\n=== ReduceScatter Results ===" << std::endl; + std::cout << "Rank " << rank << " owns block " << rank << std::endl; + std::cout << "Correct: " + << (correct ? (green + std::string("YES") + reset) + : (red + std::string("NO") + reset)) + << std::endl; + // Sample the first element of the received block; the rest is validated + // above. + std::cout << "Sample: " << output.front() << " (expected " + << ExpectedValue(size, static_cast(rank)) << ")" + << std::endl; +} + +void PrintMetrics(double elapsed_ms, size_t send_bytes, int size) { + const double seconds = elapsed_ms / 1000.0; + const double gigabytes = + static_cast(send_bytes) / (1024.0 * 1024.0 * 1024.0); + const double alg_bw = gigabytes / seconds; + // `ReduceScatter` moves `(n - 1) / n` of the send buffer per rank. + const double bus_bw = + (static_cast(size - 1) / static_cast(size)) * alg_bw; + + std::cout << "Time: " << std::fixed << std::setprecision(3) + << elapsed_ms << " ms" << std::endl; + std::cout << "Throughput: " << std::fixed << std::setprecision(2) + << bus_bw << " GB/s (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: " << std::fixed << std::setprecision(2) + << alg_bw << " GB/s" << std::endl; + return; +} + +void WorkerThread(ThreadArgs args) { + constexpr Device::Type kDevType = + ListGetBest(EnabledDevices{}); + using Rt = Runtime; + + // Bind this specific CPU thread to its designated local GPU device. + CHECK_RT(Rt, Rt::SetDevice(args.rank)); + + infinicclComm_t comm = nullptr; + CHECK_INFINI(infinicclCommInitRank(&comm, args.size, args.id, args.rank)); + + // The send buffer holds one block per rank, the receive buffer one block. + const size_t recv_bytes = args.recv_count * sizeof(float); + const size_t send_bytes = recv_bytes * static_cast(args.size); + + std::vector host_send(send_bytes / sizeof(float)); + std::vector host_recv(args.recv_count, 0.0f); + for (size_t block = 0; block < static_cast(args.size); ++block) { + std::fill_n(host_send.begin() + block * args.recv_count, args.recv_count, + SeedValue(args.rank, block)); + } + + // Allocate the device buffers once and keep the data resident for the whole + // run: the timed loop never touches the host. + float* device_send = nullptr; + float* device_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&device_send), send_bytes)); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&device_recv), recv_bytes)); + CHECK_RT(Rt, Rt::Memcpy(device_send, host_send.data(), send_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + auto reduce_scatter = [&]() { + return infinicclReduceScatter(device_send, device_recv, args.recv_count, + infinicclFloat32, infinicclSum, comm, + nullptr); + }; + + for (int i = 0; i < args.warmup_iter; ++i) { + CHECK_INFINI(reduce_scatter()); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + // All launches are enqueued back to back and the stream is synchronized only + // once, so the measurement reflects the collective itself rather than the + // per-iteration launch latency. + Timer timer; + for (int i = 0; i < args.profile_iter; ++i) { + CHECK_INFINI(reduce_scatter()); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const double elapsed = + timer.ElapsedMs() / static_cast(args.profile_iter); + + CHECK_RT(Rt, Rt::Memcpy(host_recv.data(), device_recv, recv_bytes, + Rt::MemcpyDeviceToHost)); + + const bool correct = Validate(host_recv, args.size, args.rank); + if (!correct) { + args.all_correct->store(false, std::memory_order_relaxed); + } + + if (args.rank == 0) { + PrintResult(host_recv, args.rank, args.size, correct); + std::cout << "\n=== Single-Node Threaded ReduceScatter Results ===" + << std::endl; + std::cout << "Receive size: " << args.recv_count << " floats per rank (" + << recv_bytes / 1024 / 1024 + << " MB), send size: " << send_bytes / 1024 / 1024 << " MB" + << std::endl; + PrintMetrics(elapsed, send_bytes, args.size); + } + + // Cleanup rank-local resources. + CHECK_RT(Rt, Rt::Free(device_send)); + CHECK_RT(Rt, Rt::Free(device_recv)); + CHECK_INFINI(infinicclCommDestroy(comm)); +} + +} // 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; + while ((opt = getopt(argc, argv, "g:w:p:n:h")) != -1) { + switch (opt) { + case 'g': + num_gpus = std::stoi(optarg); + break; + case 'w': + warmup_iters = std::stoi(optarg); + break; + case 'p': + profile_iters = std::stoi(optarg); + break; + case 'n': + recv_count = static_cast(std::stoull(optarg)); + break; + case 'h': + std::cout << "Usage: " << argv[0] << " [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: " + << (1 << 20) << ")\n"; + return EXIT_SUCCESS; + default: + std::cerr << "Invalid argument. Use -h for help." << std::endl; + return EXIT_FAILURE; + } + } + + char hostname[256]; + gethostname(hostname, sizeof(hostname)); + std::cout << "[Main Process] Host: " << hostname + << " | Target GPUs: " << num_gpus << std::endl; + + infinicclUniqueId shared_id; + CHECK_INFINI(infinicclGetUniqueId(&shared_id)); + + std::atomic_bool all_correct{true}; + 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, &all_correct}; + threads.emplace_back(WorkerThread, args); + } + for (auto& thread : threads) { + thread.join(); + } + + if (!all_correct.load(std::memory_order_relaxed)) { + std::cerr << "ReduceScatter validation failed." << std::endl; + return EXIT_FAILURE; + } + + std::cout << "ReduceScatter validation passed." << std::endl; + return EXIT_SUCCESS; +} 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..ac01863 --- /dev/null +++ b/src/backends/ccl/common/impl/reduce_scatter.h @@ -0,0 +1,48 @@ +#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 +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; + + if (!comm || !comm->intra_comm() || comm->intra_comm_backend() != backend || + comm->device_type() != device) { + return ReturnStatus::kInternalError; + } + auto* intra = static_cast(comm->intra_comm()); + if (!intra->handle) { + return ReturnStatus::kInternalError; + } + + typename Api::DataType ccl_type{}; + if (!TypeMap::ToBackendDataType(data_type, &ccl_type)) { + return ReturnStatus::kNotSupported; + } + typename Api::RedOp ccl_op{}; + if (!TypeMap::ToBackendRedOp(op, &ccl_op)) { + return ReturnStatus::kNotSupported; + } + + return Api::Check(Api::ReduceScatter( + send_buff, recv_buff, recv_count, ccl_type, ccl_op, intra->handle, + reinterpret_cast(stream))); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_REDUCE_SCATTER_H_ diff --git a/src/backends/ccl/nccl/api.h b/src/backends/ccl/nccl/api.h index e1d51a3..75a10df 100644 --- a/src/backends/ccl/nccl/api.h +++ b/src/backends/ccl/nccl/api.h @@ -56,6 +56,13 @@ struct NcclApi { 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); + } + 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); 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_