diff --git a/examples/ccl/all_gather.cc b/examples/ccl/all_gather.cc new file mode 100644 index 0000000..9f640d7 --- /dev/null +++ b/examples/ccl/all_gather.cc @@ -0,0 +1,194 @@ +/** + * InfiniCCL Example: Thread-per-GPU Single-Node AllGather + * + * This example validates out-of-place and in-place AllGather across two GPUs + * through InfiniCCL's native CCL backend without an MPI launcher. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "backend_manifest.h" +#include "infiniccl.h" +#include "utils.h" + +using namespace infini::ccl; + +namespace { + +constexpr int kRankCount = 2; + +struct ThreadArgs { + int rank; + infinicclUniqueId id; + size_t num_elements; + int warmup_iter; + int profile_iter; + std::atomic_bool* all_correct; +}; + +bool Validate(const std::vector& output, int rank) { + // Every rank checks the complete gathered output by source-rank block. + const size_t num_elements = output.size() / kRankCount; + bool correct = true; + for (int source = 0; source < kRankCount; ++source) { + const float expected = static_cast(source + 1); + const size_t offset = static_cast(source) * num_elements; + + correct = Validator::ValidateResult(output.data() + offset, num_elements, + expected, rank) && + correct; + } + return correct; +} + +void PrintResult(const std::vector& output, const char* mode, + bool correct) { + const char* green = "\033[32m"; + const char* red = "\033[31m"; + const char* reset = "\033[0m"; + const size_t num_elements = output.size() / kRankCount; + + std::cout << "\n=== " << mode << " AllGather Results ===" << std::endl; + std::cout << "Correct: " + << (correct ? (green + std::string("YES") + reset) + : (red + std::string("NO") + reset)) + << std::endl; + std::cout << "Sample blocks: "; + for (int source = 0; source < kRankCount; ++source) { + const size_t offset = static_cast(source) * num_elements; + std::cout << "[r" << source << ": " << output[offset] << "] "; + } + std::cout << std::endl; +} + +void WorkerThread(ThreadArgs args) { + constexpr Device::Type kDevType = + ListGetBest(EnabledDevices{}); + using Rt = Runtime; + + CHECK_RT(Rt, Rt::SetDevice(args.rank)); + + infinicclComm_t comm = nullptr; + CHECK_INFINI(infinicclCommInitRank(&comm, kRankCount, args.id, args.rank)); + + std::vector host_send(args.num_elements, + static_cast(args.rank + 1)); + std::vector host_recv(args.num_elements * kRankCount, 0.0f); + + // Prepare separate send and receive buffers for the out-of-place case. + float* device_send = nullptr; + float* device_recv = nullptr; + const size_t send_bytes = args.num_elements * sizeof(float); + const size_t recv_bytes = send_bytes * kRankCount; + + 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 all_gather = [&]() { + return infinicclAllGather(device_send, device_recv, args.num_elements, + infinicclFloat32, comm, nullptr); + }; + + for (int i = 0; i < args.warmup_iter; ++i) { + CHECK_INFINI(all_gather()); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + Timer timer; + for (int i = 0; i < args.profile_iter; ++i) { + CHECK_INFINI(all_gather()); + } + 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 out_of_place_correct = Validate(host_recv, args.rank); + if (!out_of_place_correct) { + args.all_correct->store(false, std::memory_order_relaxed); + } + if (args.rank == 0) { + PrintResult(host_recv, "Out-of-place", out_of_place_correct); + } + + // Seed only the local block, then use it as both input and output. + std::fill(host_recv.begin(), host_recv.end(), 0.0f); + std::fill_n( + host_recv.begin() + static_cast(args.rank) * args.num_elements, + args.num_elements, static_cast(args.rank + 1)); + CHECK_RT(Rt, Rt::Memcpy(device_recv, host_recv.data(), recv_bytes, + Rt::MemcpyHostToDevice)); + + float* local_block = + device_recv + static_cast(args.rank) * args.num_elements; + CHECK_INFINI(infinicclAllGather(local_block, device_recv, args.num_elements, + infinicclFloat32, comm, nullptr)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + CHECK_RT(Rt, Rt::Memcpy(host_recv.data(), device_recv, recv_bytes, + Rt::MemcpyDeviceToHost)); + const bool in_place_correct = Validate(host_recv, args.rank); + if (!in_place_correct) { + args.all_correct->store(false, std::memory_order_relaxed); + } + + if (args.rank == 0) { + PrintResult(host_recv, "In-place", in_place_correct); + std::cout << "\n=== Single-Node Threaded AllGather Results ===" + << std::endl; + Metrics metrics{elapsed, recv_bytes, kRankCount}; + metrics.Print(); + } + + // 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() { + constexpr size_t kNumElements = 1 << 20; + constexpr int kWarmupIterations = 2; + constexpr int kProfileIterations = 20; + + infinicclUniqueId shared_id; + CHECK_INFINI(infinicclGetUniqueId(&shared_id)); + + std::atomic_bool all_correct{true}; + std::vector threads; + threads.reserve(kRankCount); + + for (int rank = 0; rank < kRankCount; ++rank) { + ThreadArgs args{rank, + shared_id, + kNumElements, + kWarmupIterations, + kProfileIterations, + &all_correct}; + threads.emplace_back(WorkerThread, args); + } + for (auto& thread : threads) { + thread.join(); + } + + if (!all_correct.load(std::memory_order_relaxed)) { + std::cerr << "AllGather validation failed." << std::endl; + return EXIT_FAILURE; + } + + std::cout << "AllGather validation passed." << std::endl; + return EXIT_SUCCESS; +} diff --git a/examples/ccl_mpi_hybrid/all_gather.cc b/examples/ccl_mpi_hybrid/all_gather.cc new file mode 100644 index 0000000..aa2d5be --- /dev/null +++ b/examples/ccl_mpi_hybrid/all_gather.cc @@ -0,0 +1,183 @@ +/** + * InfiniCCL Example: AllGather (MPI + CCL Hybrid) + * + * This example performs a collective gather across global ranks. MPI + * bootstraps the processes and the CCL unique ID; the collective itself uses + * the CCL backend. + */ + +#include + +#include +#include +#include +#include +#include + +#include "backend_manifest.h" +#include "device.h" +#include "infiniccl.h" +#include "runtime.h" +#include "traits.h" +#include "utils.h" + +namespace ccl = infini::ccl; + +bool RunAllGatherExample(int argc, char **argv, int warmup_iter, + int profile_iter, size_t num_elements) { + constexpr ccl::Device::Type kDevType = + ccl::ListGetBest(ccl::EnabledDevices{}); + using Rt = ccl::Runtime; + + // Initialize InfiniCCL and obtain the global rank information. + CHECK_INFINI(infinicclInit(&argc, &argv)); + + int rank = 0; + int size = 0; + CHECK_INFINI(infinicclGetRank(&rank)); + CHECK_INFINI(infinicclGetSize(&size)); + + char hostname[256]; + gethostname(hostname, sizeof(hostname)); + + // Map local rank to GPU device. + 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); + } + + CHECK_RT(Rt, Rt::SetDevice(local_rank)); + + // Setup the MPI-backed communicator used to bootstrap the CCL unique ID. + infinicclComm_t comm = nullptr; + CHECK_INFINI(infinicclCommInitAll(&comm, size, nullptr)); + + infinicclUniqueId id; + if (rank == 0) { + CHECK_INFINI(infinicclGetUniqueId(&id)); + } + CHECK_INFINI(infinicclBroadcast(&id, &id, sizeof(id), infinicclChar, 0, comm, + nullptr)); + + std::cout << "[Rank " << rank << "] Host: " << hostname + << " | GPU: " << ccl::Device::StringFromType(kDevType) << " " + << " | Device " << local_rank << std::endl; + + CHECK_INFINI(infinicclCommInitRank(&comm, size, id, rank)); + + // Prepare Data + std::vector h_send(num_elements, static_cast(rank + 1)); + std::vector h_recv(static_cast(size) * num_elements, 0.0f); + + float *d_send = nullptr; + float *d_recv = nullptr; + const size_t send_bytes = num_elements * sizeof(*d_send); + const size_t recv_bytes = send_bytes * static_cast(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)); + + if (rank == 0) { + std::cout << "\n=== Performing AllGather on GPU Memory ===" << std::endl; + std::cout << "Data size: " << num_elements << " floats (" + << send_bytes / 1024 / 1024 << " MB per rank)" << std::endl; + std::cout << "Operation: AllGather" << 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)); + + auto all_gather_call = [&]() { + return infinicclAllGather(d_send, d_recv, num_elements, infinicclFloat32, + comm, nullptr); + }; + + // Warm-up and D2H transfer the answer. + CHECK_INFINI(all_gather_call()); + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, recv_bytes, + Rt::MemcpyDeviceToHost)); + + for (int i = 1; i < warmup_iter; ++i) { + CHECK_INFINI(all_gather_call()); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + // Profiling + Timer timer; + + for (int i = 0; i < profile_iter; ++i) { + CHECK_INFINI(all_gather_call()); + } + + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, recv_bytes, + Rt::MemcpyDeviceToHost)); + + const double elapsed = timer.ElapsedMs() / static_cast(profile_iter); + + // Result Validation + bool correct = true; + for (int source_rank = 0; source_rank < size; ++source_rank) { + const float expected = static_cast(source_rank + 1); + const size_t offset = static_cast(source_rank) * num_elements; + + correct = Validator::ValidateResult(h_recv.data() + offset, num_elements, + expected, rank) && + correct; + } + + if (rank == 0) { + const char *kGreen = "\033[32m"; + const char *kRed = "\033[31m"; + const char *kReset = "\033[0m"; + + std::cout << "\n=== AllGather Results ===" << std::endl; + std::cout << "Correct: " + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) + << std::endl; + + std::cout << "Sample blocks: "; + for (int source_rank = 0; source_rank < std::min(size, 4); ++source_rank) { + const size_t offset = static_cast(source_rank) * num_elements; + std::cout << "[r" << source_rank << ": " << h_recv[offset] << "] "; + } + std::cout << std::endl; + } + + // Metrics Reporting (Only from rank 0 for cleaner output) + if (rank == 0) { + Metrics metrics{elapsed, recv_bytes, size}; + metrics.Print(); + } + + // Cleanup + CHECK_RT(Rt, Rt::Free(d_send)); + CHECK_RT(Rt, Rt::Free(d_recv)); + + CHECK_INFINI(infinicclCommDestroy(comm)); + CHECK_INFINI(infinicclFinalize()); + + if (rank == 0) { + std::cout << "InfiniCCL finalized." << std::endl; + } + + return correct; +} + +int main(int argc, char **argv) { + const int warmup_iters = 2; + const int profile_iters = 20; + const size_t num_elements = 1 << 20; + + return RunAllGatherExample(argc, argv, warmup_iters, profile_iters, + num_elements) + ? EXIT_SUCCESS + : EXIT_FAILURE; +} diff --git a/examples/mpi/all_gather.cc b/examples/mpi/all_gather.cc index c0f1a62..5716b20 100644 --- a/examples/mpi/all_gather.cc +++ b/examples/mpi/all_gather.cc @@ -24,7 +24,7 @@ using namespace infini::ccl; -void RunAllGatherExample(int argc, char **argv, int warmup_iter, +bool RunAllGatherExample(int argc, char **argv, int warmup_iter, int profile_iter, const size_t kNumElements) { constexpr Device::Type kDevType = ListGetBest(EnabledDevices{}); @@ -113,14 +113,14 @@ void RunAllGatherExample(int argc, char **argv, int warmup_iter, // Result Validation bool correct = true; - int error_count = 0; for (int src_rank = 0; src_rank < size; ++src_rank) { float expected = static_cast(src_rank + 1); size_t offset = static_cast(src_rank) * kNumElements; - Validator::ValidateResult(h_recv.data() + offset, kNumElements, expected, - rank); + correct = Validator::ValidateResult(h_recv.data() + offset, kNumElements, + expected, rank) && + correct; } if (rank == 0) { @@ -132,7 +132,6 @@ void RunAllGatherExample(int argc, char **argv, int warmup_iter, std::cout << "Correct: " << (correct ? (GREEN + std::string("YES") + RESET) : (RED + std::string("NO") + RESET)); - if (!correct) std::cout << " (" << error_count << " errors)"; std::cout << std::endl; std::cout << "Sample blocks: "; @@ -159,6 +158,7 @@ void RunAllGatherExample(int argc, char **argv, int warmup_iter, if (rank == 0) { std::cout << "InfiniCCL finalized." << std::endl; } + return correct; } int main(int argc, char **argv) { @@ -166,7 +166,8 @@ int main(int argc, char **argv) { int profile_iters = 20; size_t num_elements = 1 << 20; - RunAllGatherExample(argc, argv, warmup_iters, profile_iters, num_elements); - - return EXIT_SUCCESS; + return RunAllGatherExample(argc, argv, warmup_iters, profile_iters, + num_elements) + ? EXIT_SUCCESS + : EXIT_FAILURE; } diff --git a/src/backends/ccl/common/impl/all_gather.h b/src/backends/ccl/common/impl/all_gather.h new file mode 100644 index 0000000..10fc340 --- /dev/null +++ b/src/backends/ccl/common/impl/all_gather.h @@ -0,0 +1,43 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_GATHER_H_ +#define INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_GATHER_H_ + +#include "backends/ccl/common/api.h" +#include "backends/ccl/common/comm_instance.h" +#include "base/all_gather.h" +#include "communicator.h" + +namespace infini::ccl { + +template +class CclAllGatherImpl { + public: + static ReturnStatus Apply(const void* send_buff, void* recv_buff, + size_t count, DataType data_type, + Communicator* comm, void* stream) { + using Api = CclApi; + using TypeMap = CclTypeMap; + using CommInstance = CclCommInstance; + + if (!comm) { + return ReturnStatus::kInternalError; + } + + auto* intra = static_cast(comm->intra_comm()); + if (!intra || !intra->handle) { + return ReturnStatus::kInternalError; + } + + typename Api::DataType ccl_type{}; + if (!TypeMap::ToBackendDataType(data_type, &ccl_type)) { + return ReturnStatus::kNotSupported; + } + + return Api::Check( + Api::AllGather(send_buff, recv_buff, count, ccl_type, intra->handle, + reinterpret_cast(stream))); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_GATHER_H_ diff --git a/src/backends/ccl/nccl/api.h b/src/backends/ccl/nccl/api.h index ed6e62f..e1d51a3 100644 --- a/src/backends/ccl/nccl/api.h +++ b/src/backends/ccl/nccl/api.h @@ -49,6 +49,13 @@ struct NcclApi { stream); } + static Result AllGather(const void* send_buff, void* recv_buff, + size_t send_count, DataType data_type, Comm comm, + Stream stream) { + return ncclAllGather(send_buff, recv_buff, send_count, data_type, 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/all_gather.h b/src/backends/ccl/nccl/impl/all_gather.h new file mode 100644 index 0000000..fdabdd7 --- /dev/null +++ b/src/backends/ccl/nccl/impl/all_gather.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_GATHER_H_ +#define INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_GATHER_H_ + +#include "backends/ccl/common/impl/all_gather.h" + +namespace infini::ccl { + +template +class AllGatherImpl + : public CclAllGatherImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_GATHER_H_ diff --git a/src/base/all_gather.h b/src/base/all_gather.h index c0072d8..1341cf4 100644 --- a/src/base/all_gather.h +++ b/src/base/all_gather.h @@ -19,10 +19,16 @@ class AllGather : public Operation { static ReturnStatus Execute(const void *send_buff, void *recv_buff, size_t count, DataType datatype, void *comm_handle, void *stream) { + if (count == 0) { + return ReturnStatus::kSuccess; + } + if (HasInvalidArgs(send_buff, recv_buff, datatype, comm_handle)) { return ReturnStatus::kInvalidArgument; } + auto *comm = static_cast(comm_handle); + return AllGatherImpl::Apply( send_buff, recv_buff, count, datatype, comm, stream); }