diff --git a/examples/ccl/send_recv.cc b/examples/ccl/send_recv.cc new file mode 100644 index 0000000..85ec786 --- /dev/null +++ b/examples/ccl/send_recv.cc @@ -0,0 +1,134 @@ +/** + * InfiniCCL Example: Thread-per-GPU Single-Node Send/Recv + * + * This example transfers data from GPU 0 to GPU 1 through InfiniCCL's native + * CCL backend without an MPI launcher. + */ + +#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 int kSender = 0; +constexpr int kReceiver = 1; +constexpr float kSendValue = 7.0f; + +struct ThreadArgs { + int rank; + infinicclUniqueId id; + size_t num_elements; + int warmup_iter; + int profile_iter; + std::atomic_bool* all_correct; +}; + +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_buffer(args.num_elements, + args.rank == kSender ? kSendValue : 0.0f); + float* device_buffer = nullptr; + const size_t total_bytes = args.num_elements * sizeof(float); + + CHECK_RT(Rt, + Rt::Malloc(reinterpret_cast(&device_buffer), total_bytes)); + CHECK_RT(Rt, Rt::Memcpy(device_buffer, host_buffer.data(), total_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + auto exchange = [&]() { + if (args.rank == kSender) { + return infinicclSend(device_buffer, args.num_elements, infinicclFloat32, + kReceiver, comm, nullptr); + } + return infinicclRecv(device_buffer, args.num_elements, infinicclFloat32, + kSender, comm, nullptr); + }; + + for (int i = 0; i < args.warmup_iter; ++i) { + CHECK_INFINI(exchange()); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + Timer timer; + for (int i = 0; i < args.profile_iter; ++i) { + CHECK_INFINI(exchange()); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const double elapsed = + timer.ElapsedMs() / static_cast(args.profile_iter); + + if (args.rank == kReceiver) { + CHECK_RT(Rt, Rt::Memcpy(host_buffer.data(), device_buffer, total_bytes, + Rt::MemcpyDeviceToHost)); + const bool correct = + Validator::ValidateResult(host_buffer.data(), args.num_elements, + kSendValue, kSender, true, "Send/Recv"); + if (!correct) { + args.all_correct->store(false, std::memory_order_relaxed); + } + } else { + std::cout << "\n=== Single-Node Threaded Send/Recv Results ===" + << std::endl; + Metrics metrics{elapsed, total_bytes, kRankCount}; + metrics.Print(); + } + + CHECK_RT(Rt, Rt::Free(device_buffer)); + 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 << "Send/Recv validation failed." << std::endl; + return EXIT_FAILURE; + } + + std::cout << "Send/Recv validation passed." << std::endl; + return EXIT_SUCCESS; +} diff --git a/examples/ccl_mpi_hybrid/send_recv.cc b/examples/ccl_mpi_hybrid/send_recv.cc new file mode 100644 index 0000000..9d2e203 --- /dev/null +++ b/examples/ccl_mpi_hybrid/send_recv.cc @@ -0,0 +1,204 @@ +/** + * InfiniCCL Example: Send/Recv (Ompi + CCL Hybrid) + * + * This example performs point-to-point `infinicclSend` and `infinicclRecv` + * operations between global ranks. OpenMPI bootstraps the processes and the + * CCL unique ID; the data transfer itself uses the CCL backend. + */ + +#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; + +void RunSendRecvExample(int argc, char **argv, int warmup_iter, + int profile_iter, size_t num_elements, int sender, + int receiver, int required_ranks, float send_value) { + 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)); + + if (size < required_ranks) { + if (rank == sender) { + std::cerr << "Send/Recv example requires at least " << required_ranks + << " ranks." << std::endl; + } + + CHECK_INFINI(infinicclFinalize()); + return; + } + + // 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 == sender) { + CHECK_INFINI(infinicclGetUniqueId(&id)); + } + CHECK_INFINI(infinicclBroadcast(&id, &id, sizeof(id), infinicclChar, sender, + 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, send_value); + std::vector h_recv(num_elements, 0.0f); + + float *d_send = nullptr; + float *d_recv = nullptr; + size_t total_bytes = num_elements * 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 == sender) { + std::cout << "\n=== Performing Send/Recv on GPU Memory ===" << std::endl; + std::cout << "Sender rank: " << sender << std::endl; + std::cout << "Receiver rank: " << receiver << std::endl; + std::cout << "Data size: " << num_elements << " floats (" + << total_bytes / 1024 / 1024 << " MB)" << 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 send_recv_call = [&]() { + if (rank == sender) { + return infinicclSend(d_send, num_elements, infinicclFloat32, receiver, + comm, nullptr); + } + + if (rank == receiver) { + return infinicclRecv(d_recv, num_elements, infinicclFloat32, sender, comm, + nullptr); + } + + return infinicclSuccess; + }; + + // Warm-up and validate the first transfer. + CHECK_INFINI(send_recv_call()); + if (rank == receiver) { + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes, + Rt::MemcpyDeviceToHost)); + } + + for (int i = 1; i < warmup_iter; ++i) { + CHECK_INFINI(send_recv_call()); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + // Profiling + Timer timer; + + for (int i = 0; i < profile_iter; i++) { + CHECK_INFINI(send_recv_call()); + } + + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + if (rank == receiver) { + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes, + Rt::MemcpyDeviceToHost)); + } + + double elapsed = timer.ElapsedMs() / static_cast(profile_iter); + + // Result Validation + if (rank == receiver) { + bool correct = Validator::ValidateResult( + h_recv.data(), num_elements, send_value, rank, false, "SendRecv"); + + const char *kGreen = "\033[32m"; + const char *kRed = "\033[31m"; + const char *kReset = "\033[0m"; + + std::cout << "\n=== Send/Recv Results ===" << std::endl; + std::cout << "Correct: " + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) + << std::endl; + std::cout << "Expect: " << send_value << std::endl; + std::cout << "Actual: " << h_recv[0] << std::endl; + + if (!correct) { + CHECK_RT(Rt, Rt::Free(d_send)); + CHECK_RT(Rt, Rt::Free(d_recv)); + CHECK_INFINI(infinicclCommDestroy(comm)); + CHECK_INFINI(infinicclFinalize()); + std::exit(EXIT_FAILURE); + } + } + + // Metrics Reporting (Only from the sender for cleaner output) + if (rank == sender) { + Metrics metrics{elapsed, total_bytes, required_ranks}; + 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 == sender) { + std::cout << "InfiniCCL finalized." << std::endl; + } +} + +int main(int argc, char **argv) { + int warmup_iters = 2; + int profile_iters = 20; + size_t num_elements = 1 << 20; + constexpr int kSender = 0; + constexpr int kReceiver = 1; + constexpr int kRequiredRanks = 2; + constexpr float kSendValue = 7.0f; + + RunSendRecvExample(argc, argv, warmup_iters, profile_iters, num_elements, + kSender, kReceiver, kRequiredRanks, kSendValue); + + return EXIT_SUCCESS; +} diff --git a/src/backends/ccl/common/impl/recv.h b/src/backends/ccl/common/impl/recv.h new file mode 100644 index 0000000..b640257 --- /dev/null +++ b/src/backends/ccl/common/impl/recv.h @@ -0,0 +1,43 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_RECV_H_ +#define INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_RECV_H_ + +#include "backends/ccl/common/api.h" +#include "backends/ccl/common/comm_instance.h" +#include "base/recv.h" +#include "communicator.h" + +namespace infini::ccl { + +template +class CclRecvImpl { + public: + static ReturnStatus Apply(void *recv_buff, size_t count, DataType data_type, + int peer, Communicator *comm, void *stream) { + using Api = CclApi; + using TypeMap = CclTypeMap; + using CommInstance = CclCommInstance; + + auto *comm_internal = static_cast(comm); + if (!comm_internal) { + return ReturnStatus::kInternalError; + } + + auto *intra = static_cast(comm_internal->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::Recv(recv_buff, count, ccl_type, peer, intra->handle, + reinterpret_cast(stream))); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_RECV_H_ diff --git a/src/backends/ccl/common/impl/send.h b/src/backends/ccl/common/impl/send.h new file mode 100644 index 0000000..3dec119 --- /dev/null +++ b/src/backends/ccl/common/impl/send.h @@ -0,0 +1,44 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_SEND_H_ +#define INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_SEND_H_ + +#include "backends/ccl/common/api.h" +#include "backends/ccl/common/comm_instance.h" +#include "base/send.h" +#include "communicator.h" + +namespace infini::ccl { + +template +class CclSendImpl { + public: + static ReturnStatus Apply(const void *send_buff, size_t count, + DataType data_type, int peer, Communicator *comm, + void *stream) { + using Api = CclApi; + using TypeMap = CclTypeMap; + using CommInstance = CclCommInstance; + + auto *comm_internal = static_cast(comm); + if (!comm_internal) { + return ReturnStatus::kInternalError; + } + + auto *intra = static_cast(comm_internal->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::Send(send_buff, count, ccl_type, peer, intra->handle, + reinterpret_cast(stream))); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_SEND_H_ diff --git a/src/backends/ccl/nccl/api.h b/src/backends/ccl/nccl/api.h index dbbf24c..ed6e62f 100644 --- a/src/backends/ccl/nccl/api.h +++ b/src/backends/ccl/nccl/api.h @@ -34,20 +34,30 @@ struct NcclApi { return ReturnStatus::kSuccess; } - static Result GetUniqueId(UniqueId *id) { return ncclGetUniqueId(id); } + static Result GetUniqueId(UniqueId* id) { return ncclGetUniqueId(id); } - static Result CommInitRank(Comm *comm, int nranks, UniqueId id, int rank) { + static Result CommInitRank(Comm* comm, int nranks, UniqueId id, int rank) { return ncclCommInitRank(comm, nranks, id, rank); } static Result CommDestroy(Comm comm) { return ncclCommDestroy(comm); } - static Result AllReduce(const void *send_buff, void *recv_buff, size_t count, + static Result AllReduce(const void* send_buff, void* recv_buff, size_t count, DataType data_type, RedOp op, Comm comm, Stream stream) { return ncclAllReduce(send_buff, recv_buff, 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); + } + + static Result Recv(void* recv_buff, size_t count, DataType data_type, + int peer, Comm comm, Stream stream) { + return ncclRecv(recv_buff, count, data_type, peer, comm, stream); + } }; } // namespace infini::ccl diff --git a/src/backends/ccl/nccl/impl/recv.h b/src/backends/ccl/nccl/impl/recv.h new file mode 100644 index 0000000..dae63b5 --- /dev/null +++ b/src/backends/ccl/nccl/impl/recv.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_RECV_H_ +#define INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_RECV_H_ + +#include "backends/ccl/common/impl/recv.h" + +namespace infini::ccl { + +template +class RecvImpl + : public CclRecvImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_RECV_H_ diff --git a/src/backends/ccl/nccl/impl/send.h b/src/backends/ccl/nccl/impl/send.h new file mode 100644 index 0000000..d91cb00 --- /dev/null +++ b/src/backends/ccl/nccl/impl/send.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_SEND_H_ +#define INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_SEND_H_ + +#include "backends/ccl/common/impl/send.h" + +namespace infini::ccl { + +template +class SendImpl + : public CclSendImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_SEND_H_