From 1878506cff9e5601dfbe4451c3a1c3910f46101d Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 12 Aug 2026 04:14:17 +0800 Subject: [PATCH 1/2] feat(nccl): add point-to-point communication --- examples/ccl/send_recv.cc | 134 ++++++++++++++++++++++++++++ src/backends/ccl/common/impl/recv.h | 43 +++++++++ src/backends/ccl/common/impl/send.h | 44 +++++++++ src/backends/ccl/nccl/api.h | 16 +++- src/backends/ccl/nccl/impl/recv.h | 17 ++++ src/backends/ccl/nccl/impl/send.h | 17 ++++ src/base/recv.h | 38 ++++++-- src/base/send.h | 40 +++++++-- src/operation.h | 32 +++++++ 9 files changed, 367 insertions(+), 14 deletions(-) create mode 100644 examples/ccl/send_recv.cc create mode 100644 src/backends/ccl/common/impl/recv.h create mode 100644 src/backends/ccl/common/impl/send.h create mode 100644 src/backends/ccl/nccl/impl/recv.h create mode 100644 src/backends/ccl/nccl/impl/send.h 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/src/backends/ccl/common/impl/recv.h b/src/backends/ccl/common/impl/recv.h new file mode 100644 index 0000000..ef5ab39 --- /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; + + 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::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..eea4f12 --- /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; + + 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::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_ diff --git a/src/base/recv.h b/src/base/recv.h index 43d89c1..14797ea 100644 --- a/src/base/recv.h +++ b/src/base/recv.h @@ -15,14 +15,14 @@ struct RecvImpl; class Recv : public Operation { public: template - static ReturnStatus Execute(void *recv_buff, size_t count, DataType datatype, - int peer, void *comm_handle, void *stream) { + static ReturnStatus Execute(void* recv_buff, size_t count, DataType datatype, + int peer, void* comm_handle, void* stream) { if (!comm_handle) { LOG("Invalid communicator handle for `Recv`."); return ReturnStatus::kInvalidArgument; } - auto *comm = static_cast(comm_handle); + auto* comm = static_cast(comm_handle); if (HasInvalidArgs(recv_buff, count, datatype, peer, comm)) { return ReturnStatus::kInvalidArgument; } @@ -30,13 +30,41 @@ class Recv : public Operation { 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 `Recv`."); + return ReturnStatus::kInternalError; + } + return ReturnStatus::kNotSupported; + } + + return Operation::Call(comm_backend, comm->device_type(), + recv_buff, count, datatype, peer, + comm_handle, stream); + } + return RecvImpl::Apply( recv_buff, count, datatype, peer, comm, stream); } private: - static bool HasInvalidArgs(const void *recv_buff, size_t count, - DataType datatype, int peer, Communicator *comm) { + template + struct BackendDependentType { + using type = T; + }; + + static bool HasInvalidArgs(const void* recv_buff, size_t count, + DataType datatype, int peer, Communicator* comm) { if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) { LOG("Invalid data type for `Recv`."); return true; diff --git a/src/base/send.h b/src/base/send.h index fbc3657..718c034 100644 --- a/src/base/send.h +++ b/src/base/send.h @@ -15,15 +15,15 @@ struct SendImpl; class Send : public Operation { public: template - static ReturnStatus Execute(const void *send_buff, size_t count, - DataType datatype, int peer, void *comm_handle, - void *stream) { + static ReturnStatus Execute(const void* send_buff, size_t count, + DataType datatype, int peer, void* comm_handle, + void* stream) { if (!comm_handle) { LOG("Invalid communicator handle for `Send`."); return ReturnStatus::kInvalidArgument; } - auto *comm = static_cast(comm_handle); + auto* comm = static_cast(comm_handle); if (HasInvalidArgs(send_buff, count, datatype, peer, comm)) { return ReturnStatus::kInvalidArgument; } @@ -31,13 +31,41 @@ class Send : public Operation { 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 `Send`."); + return ReturnStatus::kInternalError; + } + return ReturnStatus::kNotSupported; + } + + return Operation::Call(comm_backend, comm->device_type(), + send_buff, count, datatype, peer, + comm_handle, stream); + } + return SendImpl::Apply( send_buff, count, datatype, peer, comm, stream); } private: - static bool HasInvalidArgs(const void *send_buff, size_t count, - DataType datatype, int peer, Communicator *comm) { + template + struct BackendDependentType { + using type = T; + }; + + static bool HasInvalidArgs(const void* send_buff, size_t count, + DataType datatype, int peer, Communicator* comm) { if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) { LOG("Invalid data type for `Send`."); return true; diff --git a/src/operation.h b/src/operation.h index 60c1a38..69cec2e 100644 --- a/src/operation.h +++ b/src/operation.h @@ -1,6 +1,7 @@ #ifndef INFINI_CCL_OPERATION_H_ #define INFINI_CCL_OPERATION_H_ +#include #include #include "backend.h" @@ -45,6 +46,37 @@ class Operation { }, "Operation::Call"); } + + static BackendType FindSupportedBackend( + Device::Type device, std::initializer_list candidates) { + for (BackendType candidate : candidates) { + if (Supports(candidate, device)) { + return candidate; + } + } + return BackendType::kCount; + } + + private: + template + static constexpr bool SupportsBackend(BackendType backend, + List) { + return ((backend == backends && + IsSupportedCombination::value) || + ...); + } + + template + static constexpr bool Supports(BackendType backend, Device::Type device, + List) { + return ((device == devices && + SupportsBackend(backend, ActiveBackends{})) || + ...); + } + + static constexpr bool Supports(BackendType backend, Device::Type device) { + return Supports(backend, device, ActiveDevices{}); + } }; } // namespace infini::ccl From 2bd221ef3065deba26ede2f1bb3a9b75fef61fb1 Mon Sep 17 00:00:00 2001 From: Zimin Li Date: Fri, 11 Sep 2026 09:23:00 +0000 Subject: [PATCH 2/2] fix: simplify point-to-point dispatch - remove communicator redispatch from the shared operation layer - align CCL `Send`/`Recv` implementations with the existing collective style - add a hybrid example for global-rank `Send`/`Recv` across MPI-launched processes --- examples/ccl_mpi_hybrid/send_recv.cc | 204 +++++++++++++++++++++++++++ src/backends/ccl/common/impl/recv.h | 12 +- src/backends/ccl/common/impl/send.h | 14 +- src/base/recv.h | 38 +---- src/base/send.h | 40 +----- src/operation.h | 32 ----- 6 files changed, 228 insertions(+), 112 deletions(-) create mode 100644 examples/ccl_mpi_hybrid/send_recv.cc 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 index ef5ab39..b640257 100644 --- a/src/backends/ccl/common/impl/recv.h +++ b/src/backends/ccl/common/impl/recv.h @@ -11,19 +11,19 @@ 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) { + 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; - if (!comm || !comm->intra_comm() || comm->intra_comm_backend() != backend || - comm->device_type() != device) { + auto *comm_internal = static_cast(comm); + if (!comm_internal) { return ReturnStatus::kInternalError; } - auto* intra = static_cast(comm->intra_comm()); - if (!intra->handle) { + auto *intra = static_cast(comm_internal->intra_comm()); + if (!intra || !intra->handle) { return ReturnStatus::kInternalError; } diff --git a/src/backends/ccl/common/impl/send.h b/src/backends/ccl/common/impl/send.h index eea4f12..3dec119 100644 --- a/src/backends/ccl/common/impl/send.h +++ b/src/backends/ccl/common/impl/send.h @@ -11,20 +11,20 @@ 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) { + 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; - if (!comm || !comm->intra_comm() || comm->intra_comm_backend() != backend || - comm->device_type() != device) { + auto *comm_internal = static_cast(comm); + if (!comm_internal) { return ReturnStatus::kInternalError; } - auto* intra = static_cast(comm->intra_comm()); - if (!intra->handle) { + auto *intra = static_cast(comm_internal->intra_comm()); + if (!intra || !intra->handle) { return ReturnStatus::kInternalError; } diff --git a/src/base/recv.h b/src/base/recv.h index 14797ea..43d89c1 100644 --- a/src/base/recv.h +++ b/src/base/recv.h @@ -15,14 +15,14 @@ struct RecvImpl; class Recv : public Operation { public: template - static ReturnStatus Execute(void* recv_buff, size_t count, DataType datatype, - int peer, void* comm_handle, void* stream) { + static ReturnStatus Execute(void *recv_buff, size_t count, DataType datatype, + int peer, void *comm_handle, void *stream) { if (!comm_handle) { LOG("Invalid communicator handle for `Recv`."); return ReturnStatus::kInvalidArgument; } - auto* comm = static_cast(comm_handle); + auto *comm = static_cast(comm_handle); if (HasInvalidArgs(recv_buff, count, datatype, peer, comm)) { return ReturnStatus::kInvalidArgument; } @@ -30,41 +30,13 @@ class Recv : public Operation { 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 `Recv`."); - return ReturnStatus::kInternalError; - } - return ReturnStatus::kNotSupported; - } - - return Operation::Call(comm_backend, comm->device_type(), - recv_buff, count, datatype, peer, - comm_handle, stream); - } - return RecvImpl::Apply( recv_buff, count, datatype, peer, comm, stream); } private: - template - struct BackendDependentType { - using type = T; - }; - - static bool HasInvalidArgs(const void* recv_buff, size_t count, - DataType datatype, int peer, Communicator* comm) { + static bool HasInvalidArgs(const void *recv_buff, size_t count, + DataType datatype, int peer, Communicator *comm) { if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) { LOG("Invalid data type for `Recv`."); return true; diff --git a/src/base/send.h b/src/base/send.h index 718c034..fbc3657 100644 --- a/src/base/send.h +++ b/src/base/send.h @@ -15,15 +15,15 @@ struct SendImpl; class Send : public Operation { public: template - static ReturnStatus Execute(const void* send_buff, size_t count, - DataType datatype, int peer, void* comm_handle, - void* stream) { + static ReturnStatus Execute(const void *send_buff, size_t count, + DataType datatype, int peer, void *comm_handle, + void *stream) { if (!comm_handle) { LOG("Invalid communicator handle for `Send`."); return ReturnStatus::kInvalidArgument; } - auto* comm = static_cast(comm_handle); + auto *comm = static_cast(comm_handle); if (HasInvalidArgs(send_buff, count, datatype, peer, comm)) { return ReturnStatus::kInvalidArgument; } @@ -31,41 +31,13 @@ class Send : public Operation { 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 `Send`."); - return ReturnStatus::kInternalError; - } - return ReturnStatus::kNotSupported; - } - - return Operation::Call(comm_backend, comm->device_type(), - send_buff, count, datatype, peer, - comm_handle, stream); - } - return SendImpl::Apply( send_buff, count, datatype, peer, comm, stream); } private: - template - struct BackendDependentType { - using type = T; - }; - - static bool HasInvalidArgs(const void* send_buff, size_t count, - DataType datatype, int peer, Communicator* comm) { + static bool HasInvalidArgs(const void *send_buff, size_t count, + DataType datatype, int peer, Communicator *comm) { if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) { LOG("Invalid data type for `Send`."); return true; diff --git a/src/operation.h b/src/operation.h index 69cec2e..60c1a38 100644 --- a/src/operation.h +++ b/src/operation.h @@ -1,7 +1,6 @@ #ifndef INFINI_CCL_OPERATION_H_ #define INFINI_CCL_OPERATION_H_ -#include #include #include "backend.h" @@ -46,37 +45,6 @@ class Operation { }, "Operation::Call"); } - - static BackendType FindSupportedBackend( - Device::Type device, std::initializer_list candidates) { - for (BackendType candidate : candidates) { - if (Supports(candidate, device)) { - return candidate; - } - } - return BackendType::kCount; - } - - private: - template - static constexpr bool SupportsBackend(BackendType backend, - List) { - return ((backend == backends && - IsSupportedCombination::value) || - ...); - } - - template - static constexpr bool Supports(BackendType backend, Device::Type device, - List) { - return ((device == devices && - SupportsBackend(backend, ActiveBackends{})) || - ...); - } - - static constexpr bool Supports(BackendType backend, Device::Type device) { - return Supports(backend, device, ActiveDevices{}); - } }; } // namespace infini::ccl