From bc387666d5a65e1f8d35986ff3efd9edc6e430fd Mon Sep 17 00:00:00 2001 From: lqinfdim <183612562+lqinfdim@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:20:48 +0800 Subject: [PATCH] feat: add point-to-point communication for mccl --- CMakeLists.txt | 5 + examples/ccl/send_recv.cc | 194 ++++++++++++++++++++++ src/CMakeLists.txt | 17 ++ src/backends/ccl/common/impl/recv.h | 43 +++++ src/backends/ccl/common/impl/send.h | 44 +++++ src/backends/ccl/mccl/api.h | 16 +- src/backends/ccl/mccl/impl/recv.h | 17 ++ src/backends/ccl/mccl/impl/send.h | 17 ++ tests/CMakeLists.txt | 26 +++ tests/bridge_dependency_contract.py | 44 +++++ tests/bridge_generation.py | 48 ++++++ tests/ccl_point_to_point_impl.cc | 244 ++++++++++++++++++++++++++++ 12 files changed, 712 insertions(+), 3 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/mccl/impl/recv.h create mode 100644 src/backends/ccl/mccl/impl/send.h create mode 100644 tests/CMakeLists.txt create mode 100644 tests/bridge_dependency_contract.py create mode 100644 tests/bridge_generation.py create mode 100644 tests/ccl_point_to_point_impl.cc diff --git a/CMakeLists.txt b/CMakeLists.txt index ebaca2c..565f79a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,6 +1,8 @@ cmake_minimum_required(VERSION 3.18) project(InfiniCCL VERSION 0.1.0 LANGUAGES C CXX) +include(CTest) + set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) @@ -507,6 +509,9 @@ add_subdirectory(src) if(BUILD_EXAMPLES) add_subdirectory(examples) endif() +if(BUILD_TESTING) + add_subdirectory(tests) +endif() # ========================================================= # --- Installation --- diff --git a/examples/ccl/send_recv.cc b/examples/ccl/send_recv.cc new file mode 100644 index 0000000..6bd2815 --- /dev/null +++ b/examples/ccl/send_recv.cc @@ -0,0 +1,194 @@ +/** + * 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 +#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; +}; + +template +bool ParseNumber(const char* text, T* value) { + const std::string_view input{text}; + T parsed{}; + const auto [end, error] = + std::from_chars(input.data(), input.data() + input.size(), parsed); + + if (error != std::errc{} || end != input.data() + input.size()) { + return false; + } + + *value = parsed; + 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_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(int argc, char** argv) { + size_t num_elements = 1 << 20; + int warmup_iterations = 2; + int profile_iterations = 20; + + int opt; + while ((opt = getopt(argc, argv, "n:w:p:h")) != -1) { + switch (opt) { + case 'n': + if (!ParseNumber(optarg, &num_elements)) { + std::cerr << "Invalid value for `-n`." << std::endl; + return EXIT_FAILURE; + } + break; + case 'w': + if (!ParseNumber(optarg, &warmup_iterations)) { + std::cerr << "Invalid value for `-w`." << std::endl; + return EXIT_FAILURE; + } + break; + case 'p': + if (!ParseNumber(optarg, &profile_iterations)) { + std::cerr << "Invalid value for `-p`." << std::endl; + return EXIT_FAILURE; + } + break; + case 'h': + std::cout << "Usage: " << argv[0] << " [options]\n" + << " -n Elements to transfer (default: " + << (1 << 20) << ")\n" + << " -w Warm-up iterations (default: 2)\n" + << " -p Profile iterations (default: 20)\n"; + return EXIT_SUCCESS; + default: + return EXIT_FAILURE; + } + } + + if (num_elements == 0 || warmup_iterations < 0 || + profile_iterations <= 0) { + std::cerr << "Elements and profile iterations must be positive; warm-up " + "iterations must be non-negative." + << std::endl; + return EXIT_FAILURE; + } + + 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, + num_elements, + warmup_iterations, + profile_iterations, + &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/CMakeLists.txt b/src/CMakeLists.txt index 1e1e726..319a546 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -16,6 +16,21 @@ set(AUTOGEN_WARNING [[/* file(GLOB CORE_SRCS "*.cc") file(GLOB_RECURSE BASE_IMPL_SRCS "base/*.cc") +file(GLOB_RECURSE BRIDGE_DEP_HEADERS CONFIGURE_DEPENDS + "base/*.h" + "backends/*.h" + "devices/*.h" + "devices/*.cuh" +) + +string(REPLACE ";" "\n" BRIDGE_DEPENDENCY_CONTENT "${BRIDGE_DEP_HEADERS}") +set(BRIDGE_DEPENDENCY_MANIFEST + "${CMAKE_CURRENT_BINARY_DIR}/bridge_dependencies.txt" +) +file(GENERATE + OUTPUT "${BRIDGE_DEPENDENCY_MANIFEST}" + CONTENT "${BRIDGE_DEPENDENCY_CONTENT}\n" +) target_sources(infiniccl PRIVATE ${CORE_SRCS} @@ -298,7 +313,9 @@ add_custom_command( "${BACK_STR}" DEPENDS "${PROJECT_SOURCE_DIR}/scripts/gen_bridge.py" "${PROJECT_SOURCE_DIR}/include/comm.h" + "${BRIDGE_DEPENDENCY_MANIFEST}" ${BASE_IMPL_SRCS} + ${BRIDGE_DEP_HEADERS} VERBATIM COMMENT "Generating InfiniCCL bridge and manifest files for Devices: [${DEV_STR}] Backends: [${BACK_STR}]..." ) 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/mccl/api.h b/src/backends/ccl/mccl/api.h index a5d2bcd..eac337d 100644 --- a/src/backends/ccl/mccl/api.h +++ b/src/backends/ccl/mccl/api.h @@ -35,20 +35,30 @@ struct McclApi { return ReturnStatus::kSuccess; } - static Result GetUniqueId(UniqueId *id) { return mcclGetUniqueId(id); } + static Result GetUniqueId(UniqueId* id) { return mcclGetUniqueId(id); } - static Result CommInitRank(Comm *comm, int nranks, UniqueId id, int rank) { + static Result CommInitRank(Comm* comm, int nranks, UniqueId id, int rank) { return mcclCommInitRank(comm, nranks, id, rank); } static Result CommDestroy(Comm comm) { return mcclCommDestroy(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 mcclAllReduce(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 mcclSend(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 mcclRecv(recv_buff, count, data_type, peer, comm, stream); + } }; } // namespace infini::ccl diff --git a/src/backends/ccl/mccl/impl/recv.h b/src/backends/ccl/mccl/impl/recv.h new file mode 100644 index 0000000..889a956 --- /dev/null +++ b/src/backends/ccl/mccl/impl/recv.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_RECV_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_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_MCCL_IMPL_RECV_H_ diff --git a/src/backends/ccl/mccl/impl/send.h b/src/backends/ccl/mccl/impl/send.h new file mode 100644 index 0000000..2c59a85 --- /dev/null +++ b/src/backends/ccl/mccl/impl/send.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_SEND_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_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_MCCL_IMPL_SEND_H_ diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt new file mode 100644 index 0000000..158a180 --- /dev/null +++ b/tests/CMakeLists.txt @@ -0,0 +1,26 @@ +add_test( + NAME bridge_dependency_contract + COMMAND ${Python3_EXECUTABLE} + "${CMAKE_CURRENT_SOURCE_DIR}/bridge_dependency_contract.py" + "${PROJECT_SOURCE_DIR}/src/CMakeLists.txt" +) + +add_test( + NAME bridge_generation + COMMAND ${Python3_EXECUTABLE} + "${CMAKE_CURRENT_SOURCE_DIR}/bridge_generation.py" + "${PROJECT_SOURCE_DIR}" +) + +add_executable(ccl_point_to_point_impl_test + ccl_point_to_point_impl.cc +) + +target_include_directories(ccl_point_to_point_impl_test PRIVATE + "${PROJECT_SOURCE_DIR}/src" +) + +add_test( + NAME ccl_point_to_point_impl + COMMAND ccl_point_to_point_impl_test +) diff --git a/tests/bridge_dependency_contract.py b/tests/bridge_dependency_contract.py new file mode 100644 index 0000000..8fdb553 --- /dev/null +++ b/tests/bridge_dependency_contract.py @@ -0,0 +1,44 @@ +import pathlib +import re +import sys + + +def main(): + source = pathlib.Path(sys.argv[1]).read_text(encoding="utf-8") + + header_glob = re.search( + r"file\(GLOB_RECURSE BRIDGE_DEP_HEADERS CONFIGURE_DEPENDS(.*?)\n\)", + source, + re.DOTALL, + ) + assert header_glob is not None + assert set(re.findall(r'"([^"]+)"', header_glob.group(1))) == { + "base/*.h", + "backends/*.h", + "devices/*.h", + "devices/*.cuh", + } + + assert re.search( + r"file\(GENERATE\s+" + r'OUTPUT "\$\{BRIDGE_DEPENDENCY_MANIFEST\}"\s+' + r'CONTENT "\$\{BRIDGE_DEPENDENCY_CONTENT\}\\n"\s*' + r"\)", + source, + ) + + bridge_command = re.search( + r"add_custom_command\(\s+" + r'OUTPUT "\$\{GENERATED_BRIDGE\}" "\$\{GENERATED_MANIFEST\}"' + r"(.*?)\n\)", + source, + re.DOTALL, + ) + assert bridge_command is not None + dependencies = bridge_command.group(1) + assert '"${BRIDGE_DEPENDENCY_MANIFEST}"' in dependencies + assert "${BRIDGE_DEP_HEADERS}" in dependencies + + +if __name__ == "__main__": + main() diff --git a/tests/bridge_generation.py b/tests/bridge_generation.py new file mode 100644 index 0000000..a34c77d --- /dev/null +++ b/tests/bridge_generation.py @@ -0,0 +1,48 @@ +import pathlib +import subprocess +import sys +import tempfile + + +def main(): + project_root = pathlib.Path(sys.argv[1]).resolve() + generator = project_root / "scripts" / "gen_bridge.py" + + with tempfile.TemporaryDirectory() as output_dir: + subprocess.run( + [ + sys.executable, + str(generator), + str(project_root), + output_dir, + "cpu;metax", + "mccl", + ], + check=True, + ) + + output = pathlib.Path(output_dir) + manifest = (output / "backend_manifest.h").read_text(encoding="utf-8") + bridge = (output / "comm_bridge.cc").read_text(encoding="utf-8") + + expected_headers = { + '#include "backends/ccl/mccl/metax/api.h"', + '#include "backends/ccl/mccl/type_map.h"', + '#include "backends/ccl/mccl/impl/send.h"', + '#include "backends/ccl/mccl/impl/recv.h"', + } + for header in expected_headers: + assert header in manifest + + send = bridge.index("infinicclResult_t infinicclSend(") + recv = bridge.index("infinicclResult_t infinicclRecv(") + assert "Operation::Call" in bridge[send:recv] + assert "ReturnStatus::kNotSupported" not in bridge[send:recv] + + recv_body = bridge[recv:] + assert "Operation::Call" in recv_body + assert "ReturnStatus::kNotSupported" not in recv_body.split("}", 1)[0] + + +if __name__ == "__main__": + main() diff --git a/tests/ccl_point_to_point_impl.cc b/tests/ccl_point_to_point_impl.cc new file mode 100644 index 0000000..9321048 --- /dev/null +++ b/tests/ccl_point_to_point_impl.cc @@ -0,0 +1,244 @@ +#include +#include +#include + +#include "backends/ccl/common/impl/recv.h" +#include "backends/ccl/common/impl/send.h" + +namespace infini::ccl { + +struct FakeMcclApi { + static constexpr BackendType kBackendType = BackendType::kMccl; + + using Comm = int; + using DataType = int; + using Result = int; + using Stream = void*; + + static constexpr Result kSuccess = 0; + static constexpr Result kFailure = 1; + + struct Call { + const void* buffer = nullptr; + size_t count = 0; + DataType data_type = 0; + int peer = -1; + Comm comm = 0; + Stream stream = nullptr; + }; + + static Call send_call; + static Call recv_call; + static Result next_result; + static int destroy_count; + + static void Reset() { + send_call = {}; + recv_call = {}; + next_result = kSuccess; + } + + static ReturnStatus Check(Result result) { + return result == kSuccess ? ReturnStatus::kSuccess + : ReturnStatus::kSystemError; + } + + static Result CommDestroy(Comm) { + ++destroy_count; + return kSuccess; + } + + static Result Send(const void* buffer, size_t count, DataType data_type, + int peer, Comm comm, Stream stream) { + send_call = {buffer, count, data_type, peer, comm, stream}; + return next_result; + } + + static Result Recv(void* buffer, size_t count, DataType data_type, int peer, + Comm comm, Stream stream) { + recv_call = {buffer, count, data_type, peer, comm, stream}; + return next_result; + } +}; + +FakeMcclApi::Call FakeMcclApi::send_call{}; +FakeMcclApi::Call FakeMcclApi::recv_call{}; +FakeMcclApi::Result FakeMcclApi::next_result = FakeMcclApi::kSuccess; +int FakeMcclApi::destroy_count = 0; + +template <> +struct CclApi : FakeMcclApi {}; + +template <> +struct CclTypeMap { + static bool ToBackendDataType(DataType data_type, int* backend_data_type) { + if (data_type == DataType::kUInt16) { + return false; + } + *backend_data_type = 100 + static_cast(data_type); + return true; + } +}; + +using Api = CclApi; +using CommInstance = CclCommInstance; +using SendUnderTest = CclSendImpl; +using RecvUnderTest = CclRecvImpl; + +std::unique_ptr MakeMcclBackend(int handle = 41) { + auto instance = std::make_unique(); + instance->handle = handle; + return instance; +} + +bool TestSendForwardsArguments() { + FakeMcclApi::Reset(); + Communicator comm(Device::Type::kMetax, 0); + comm.set_intra_comm(MakeMcclBackend()); + + float buffer = 7.0f; + void* stream = reinterpret_cast(0x1234); + auto status = SendUnderTest::Apply(&buffer, 9, DataType::kFloat32, 3, &comm, + stream); + const auto& call = FakeMcclApi::send_call; + + return status == ReturnStatus::kSuccess && call.buffer == &buffer && + call.count == 9 && + call.data_type == 100 + static_cast(DataType::kFloat32) && + call.peer == 3 && call.comm == 41 && call.stream == stream; +} + +bool TestRecvForwardsArguments() { + FakeMcclApi::Reset(); + Communicator comm(Device::Type::kMetax, 0); + comm.set_intra_comm(MakeMcclBackend(42)); + + float buffer = 0.0f; + void* stream = reinterpret_cast(0x5678); + auto status = RecvUnderTest::Apply(&buffer, 11, DataType::kInt32, 2, &comm, + stream); + const auto& call = FakeMcclApi::recv_call; + + return status == ReturnStatus::kSuccess && call.buffer == &buffer && + call.count == 11 && + call.data_type == 100 + static_cast(DataType::kInt32) && + call.peer == 2 && call.comm == 42 && call.stream == stream; +} + +bool TestUnsupportedDataTypeDoesNotCallBackend() { + FakeMcclApi::Reset(); + Communicator comm(Device::Type::kMetax, 0); + comm.set_intra_comm(MakeMcclBackend()); + + float buffer = 0.0f; + return SendUnderTest::Apply(&buffer, 1, DataType::kUInt16, 1, &comm, + nullptr) == ReturnStatus::kNotSupported && + RecvUnderTest::Apply(&buffer, 1, DataType::kUInt16, 1, &comm, + nullptr) == ReturnStatus::kNotSupported && + FakeMcclApi::send_call.buffer == nullptr && + FakeMcclApi::recv_call.buffer == nullptr; +} + +bool TestBackendErrorIsPropagated() { + FakeMcclApi::Reset(); + FakeMcclApi::next_result = FakeMcclApi::kFailure; + Communicator comm(Device::Type::kMetax, 0); + comm.set_intra_comm(MakeMcclBackend()); + + float buffer = 0.0f; + return SendUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, &comm, + nullptr) == ReturnStatus::kSystemError && + RecvUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, &comm, + nullptr) == ReturnStatus::kSystemError; +} + +bool TestInvalidCommunicatorState() { + FakeMcclApi::Reset(); + float buffer = 0.0f; + + Communicator empty(Device::Type::kMetax, 0); + if (SendUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, &empty, + nullptr) != ReturnStatus::kInternalError || + RecvUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, &empty, + nullptr) != ReturnStatus::kInternalError) { + return false; + } + + Communicator wrong_device(Device::Type::kCpu, 0); + wrong_device.set_intra_comm(MakeMcclBackend()); + if (SendUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, &wrong_device, + nullptr) != ReturnStatus::kInternalError || + RecvUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, &wrong_device, + nullptr) != ReturnStatus::kInternalError) { + return false; + } + + Communicator empty_handle(Device::Type::kMetax, 0); + empty_handle.set_intra_comm(MakeMcclBackend(0)); + return SendUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, &empty_handle, + nullptr) == ReturnStatus::kInternalError && + RecvUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, &empty_handle, + nullptr) == ReturnStatus::kInternalError; +} + +bool TestNullAndWrongBackend() { + FakeMcclApi::Reset(); + float buffer = 0.0f; + if (SendUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, nullptr, + nullptr) != ReturnStatus::kInternalError || + RecvUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, nullptr, + nullptr) != ReturnStatus::kInternalError) { + return false; + } + + auto backend = std::make_unique(); + backend->type = BackendType::kNccl; + Communicator wrong_backend(Device::Type::kMetax, 0); + wrong_backend.set_intra_comm(std::move(backend)); + return SendUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, + &wrong_backend, nullptr) == + ReturnStatus::kInternalError && + RecvUnderTest::Apply(&buffer, 1, DataType::kFloat32, 1, + &wrong_backend, nullptr) == + ReturnStatus::kInternalError; +} + +bool TestCommInstanceDestroysHandle() { + FakeMcclApi::destroy_count = 0; + { + Communicator comm(Device::Type::kMetax, 0); + comm.set_intra_comm(MakeMcclBackend()); + } + return FakeMcclApi::destroy_count == 1; +} + +bool RunTest(const char* name, bool (*test)()) { + if (test()) { + return true; + } + std::cerr << "FAILED: " << name << std::endl; + return false; +} + +} // namespace infini::ccl + +int main() { + using namespace infini::ccl; + + bool passed = true; + passed = RunTest("send forwards arguments", TestSendForwardsArguments) && + passed; + passed = RunTest("recv forwards arguments", TestRecvForwardsArguments) && + passed; + passed = RunTest("unsupported datatype", + TestUnsupportedDataTypeDoesNotCallBackend) && + passed; + passed = RunTest("backend error propagation", TestBackendErrorIsPropagated) && + passed; + passed = RunTest("invalid communicator state", TestInvalidCommunicatorState) && + passed; + passed = RunTest("null and wrong backend", TestNullAndWrongBackend) && passed; + passed = RunTest("communicator destruction", TestCommInstanceDestroysHandle) && + passed; + return passed ? EXIT_SUCCESS : EXIT_FAILURE; +}