diff --git a/examples/ccl/all_gather.cc b/examples/ccl/all_gather.cc new file mode 100644 index 0000000..e411172 --- /dev/null +++ b/examples/ccl/all_gather.cc @@ -0,0 +1,123 @@ +/** + * 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 "backend_manifest.h" +#include "infiniccl.h" +#include "utils.h" + +using namespace infini::ccl; + +namespace { + +constexpr int kRankCount = 2; +constexpr size_t kNumElements = 1 << 10; + +struct ThreadArgs { + int rank; + infinicclUniqueId id; + std::atomic_bool* all_correct; +}; + +bool Validate(const std::vector& output, int rank, const char* mode) { + for (int source = 0; source < kRankCount; ++source) { + const float expected = static_cast(source + 1); + const size_t offset = static_cast(source) * kNumElements; + for (size_t i = 0; i < kNumElements; ++i) { + if (output[offset + i] != expected) { + std::cerr << mode << " validation failed on rank " << rank + << " at source rank " << source << ", element " << i + << ": expected " << expected << ", got " << output[offset + i] + << "." << std::endl; + return false; + } + } + } + return true; +} + +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(kNumElements, static_cast(args.rank + 1)); + std::vector host_recv(kNumElements * kRankCount, 0.0f); + float* device_send = nullptr; + float* device_recv = nullptr; + const size_t send_bytes = kNumElements * 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_INFINI(infinicclAllGather(device_send, device_recv, kNumElements, + infinicclFloat32, comm, nullptr)); + CHECK_RT(Rt, Rt::Memcpy(host_recv.data(), device_recv, recv_bytes, + Rt::MemcpyDeviceToHost)); + if (!Validate(host_recv, args.rank, "Out-of-place AllGather")) { + args.all_correct->store(false, std::memory_order_relaxed); + } + + std::fill(host_recv.begin(), host_recv.end(), 0.0f); + std::fill_n(host_recv.begin() + static_cast(args.rank) * kNumElements, + kNumElements, 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) * kNumElements; + CHECK_INFINI(infinicclAllGather(local_block, device_recv, kNumElements, + infinicclFloat32, comm, nullptr)); + CHECK_RT(Rt, Rt::Memcpy(host_recv.data(), device_recv, recv_bytes, + Rt::MemcpyDeviceToHost)); + if (!Validate(host_recv, args.rank, "In-place AllGather")) { + args.all_correct->store(false, std::memory_order_relaxed); + } + + CHECK_RT(Rt, Rt::Free(device_send)); + CHECK_RT(Rt, Rt::Free(device_recv)); + CHECK_INFINI(infinicclCommDestroy(comm)); +} + +} // namespace + +int main() { + 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) { + threads.emplace_back(WorkerThread, + ThreadArgs{rank, shared_id, &all_correct}); + } + for (auto& thread : threads) { + thread.join(); + } + + if (!all_correct.load(std::memory_order_relaxed)) { + return EXIT_FAILURE; + } + std::cout << "AllGather validation passed." << std::endl; + return EXIT_SUCCESS; +} 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..ea8c95c --- /dev/null +++ b/src/backends/ccl/common/impl/all_gather.h @@ -0,0 +1,44 @@ +#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 || !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; + } + + 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 01aa099..0b2aa1f 100644 --- a/src/backends/ccl/nccl/api.h +++ b/src/backends/ccl/nccl/api.h @@ -47,6 +47,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..c7ce6ff 100644 --- a/src/base/all_gather.h +++ b/src/base/all_gather.h @@ -19,30 +19,65 @@ 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 (HasInvalidArgs(send_buff, recv_buff, datatype, comm_handle)) { + if (!comm_handle) { + LOG("Invalid communicator handle for `AllGather`."); return ReturnStatus::kInvalidArgument; } + auto *comm = static_cast(comm_handle); + if (HasInvalidArgs(send_buff, recv_buff, count, datatype)) { + return ReturnStatus::kInvalidArgument; + } + if (count == 0) { + return ReturnStatus::kSuccess; + } + + if (!comm->HasBackend(backend_type) || comm->device_type() != device_type) { + using DispatchKey = + typename BackendDependentType::type; + const BackendType comm_backend = + Operation::FindSupportedBackend( + comm->device_type(), + {comm->HasBackend(backend_type) ? backend_type + : BackendType::kCount, + comm->intra_comm_backend(), comm->inter_comm_backend()}); + if (comm_backend == BackendType::kCount) { + if (comm->intra_comm_backend() == BackendType::kCount && + comm->inter_comm_backend() == BackendType::kCount) { + LOG("No initialized backend is available for `AllGather`."); + return ReturnStatus::kInternalError; + } + return ReturnStatus::kNotSupported; + } + + return Operation::Call(comm_backend, comm->device_type(), + send_buff, recv_buff, count, datatype, + comm_handle, stream); + } + return AllGatherImpl::Apply( send_buff, recv_buff, count, datatype, comm, stream); } private: + template + struct BackendDependentType { + using type = T; + }; + static bool HasInvalidArgs(const void *send_buff, void *recv_buff, - DataType datatype, void *comm_handle) { - if (!comm_handle) { - // TODO(lzm): change to use `glog`. - LOG("Invalid communicator handle for `AllGather`."); + size_t count, DataType datatype) { + if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) { + LOG("Invalid data type for `AllGather`."); return true; } + if (count == 0) { + return false; + } if (!send_buff || !recv_buff) { LOG("Invalid buffer pointer for `AllGather`."); return true; } - if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) { - LOG("Invalid data type for `AllGather`."); - return true; - } return false; } }; diff --git a/tests/operation_backend_selection.cc b/tests/operation_backend_selection.cc index 6a6120a..50b9e73 100644 --- a/tests/operation_backend_selection.cc +++ b/tests/operation_backend_selection.cc @@ -1,6 +1,7 @@ #include #include +#include "base/all_gather.h" #include "base/recv.h" #include "base/send.h" #include "devices/cpu/device_.h" @@ -15,6 +16,17 @@ struct BackendEnabled : std::true_type {}; template <> struct BackendEnabled : std::true_type {}; +template <> +struct BackendEnabled : std::true_type {}; + +template <> +struct AllGatherImpl { + static ReturnStatus Apply(const void *, void *, size_t, DataType, + Communicator *, void *) { + return ReturnStatus::kInternalError; + } +}; + template <> struct SendImpl { static ReturnStatus Apply(const void *, size_t, DataType, int, Communicator *, @@ -47,6 +59,14 @@ struct RecvImpl { } }; +template <> +struct AllGatherImpl { + static ReturnStatus Apply(const void *, void *, size_t, DataType, + Communicator *, void *) { + return ReturnStatus::kSuccess; + } +}; + std::unique_ptr MakeBackend(BackendType backend) { auto instance = std::make_unique(); instance->type = backend; @@ -64,6 +84,9 @@ bool TestUnsupportedCommunicator() { ReturnStatus::kNotSupported && Recv::Execute( &buffer, 1, DataType::kFloat32, 1, &comm, nullptr) == + ReturnStatus::kNotSupported && + AllGather::Execute( + &buffer, &buffer, 1, DataType::kFloat32, &comm, nullptr) == ReturnStatus::kNotSupported; } @@ -84,6 +107,22 @@ bool TestInterBackendFallback() { ReturnStatus::kSuccess && Recv::Execute( &buffer, 1, DataType::kFloat32, 1, &comm, nullptr) == + ReturnStatus::kSuccess && + AllGather::Execute( + &buffer, &buffer, 1, DataType::kFloat32, &comm, nullptr) == + ReturnStatus::kSuccess; +} + +bool TestAllGatherNoBackend() { + Communicator comm(Device::Type::kCpu, 0); + comm.set_world_info(0, 2); + + float buffer = 0.0f; + return AllGather::Execute( + &buffer, &buffer, 1, DataType::kFloat32, &comm, nullptr) == + ReturnStatus::kInternalError && + AllGather::Execute( + nullptr, nullptr, 0, DataType::kFloat32, &comm, nullptr) == ReturnStatus::kSuccess; } @@ -92,7 +131,8 @@ bool TestInterBackendFallback() { int main() { if (!infini::ccl::TestNoActiveBackend() || !infini::ccl::TestUnsupportedCommunicator() || - !infini::ccl::TestInterBackendFallback()) { + !infini::ccl::TestInterBackendFallback() || + !infini::ccl::TestAllGatherNoBackend()) { return EXIT_FAILURE; } return EXIT_SUCCESS;