From 1ba15fe07801209f86858d99e27c043fb5fafc4b Mon Sep 17 00:00:00 2001 From: PanZezhong Date: Wed, 22 Jul 2026 09:19:42 +0800 Subject: [PATCH] feat: InfiniCCL initRank/Send/Recv/Broadcast --- include/infiniccl.h | 40 ++ include/infinicore/ops.hpp | 2 + .../infinicore/ops/distributed/broadcast.hpp | 25 + .../infinicore/ops/distributed/send_recv.hpp | 37 ++ src/infiniccl-test/infiniccl_test.cpp | 511 +++++++++++++++++- src/infiniccl-test/infiniccl_test.hpp | 7 + src/infiniccl-test/main.cpp | 102 +++- src/infiniccl/cuda/infiniccl_cuda.cu | 104 ++++ src/infiniccl/cuda/infiniccl_cuda.h | 70 +++ src/infiniccl/infiniccl.cc | 144 +++++ src/infinicore/ops/distributed/broadcast.cc | 54 ++ src/infinicore/ops/distributed/send_recv.cc | 88 +++ 12 files changed, 1145 insertions(+), 39 deletions(-) create mode 100644 include/infinicore/ops/distributed/broadcast.hpp create mode 100644 include/infinicore/ops/distributed/send_recv.hpp create mode 100644 src/infinicore/ops/distributed/broadcast.cc create mode 100644 src/infinicore/ops/distributed/send_recv.cc diff --git a/include/infiniccl.h b/include/infiniccl.h index b338d85f5..cced30879 100644 --- a/include/infiniccl.h +++ b/include/infiniccl.h @@ -15,12 +15,27 @@ struct InfinicclComm; typedef struct InfinicclComm *infinicclComm_t; +#define INFINICCL_UNIQUE_ID_BYTES 128 + +typedef struct { + char internal[INFINICCL_UNIQUE_ID_BYTES]; +} infinicclUniqueId_t; + __INFINI_C __export infiniStatus_t infinicclCommInitAll( infiniDevice_t device_type, infinicclComm_t *comms, int ndevice, const int *device_ids); +__INFINI_C __export infiniStatus_t infinicclGetUniqueId( + infinicclUniqueId_t *unique_id); + +__INFINI_C __export infiniStatus_t infinicclCommInitRank( + infinicclComm_t *comm, + int nranks, + infinicclUniqueId_t comm_id, + int rank); + __INFINI_C __export infiniStatus_t infinicclCommDestroy(infinicclComm_t comm); __INFINI_C __export infiniStatus_t infinicclGroupStart(infinicclComm_t comm); @@ -36,6 +51,31 @@ __INFINI_C __export infiniStatus_t infinicclAllReduce( infinicclComm_t comm, infinirtStream_t stream); +__INFINI_C __export infiniStatus_t infinicclBroadcast( + const void *sendbuf, + void *recvbuf, + size_t count, + infiniDtype_t datatype, + int root, + infinicclComm_t comm, + infinirtStream_t stream); + +__INFINI_C __export infiniStatus_t infinicclSend( + const void *sendbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream); + +__INFINI_C __export infiniStatus_t infinicclRecv( + void *recvbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream); + __INFINI_C __export infiniStatus_t infinicclAllGather( void *sendbuf, void *recvbuf, diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index b5c4ff18f..822a6b460 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -25,6 +25,8 @@ #include "ops/conv2d.hpp" #include "ops/cross_entropy.hpp" #include "ops/deepseek_moe.hpp" +#include "ops/distributed/broadcast.hpp" +#include "ops/distributed/send_recv.hpp" #include "ops/embedding.hpp" #include "ops/flash_attention.hpp" #include "ops/fmin.hpp" diff --git a/include/infinicore/ops/distributed/broadcast.hpp b/include/infinicore/ops/distributed/broadcast.hpp new file mode 100644 index 000000000..850d6f298 --- /dev/null +++ b/include/infinicore/ops/distributed/broadcast.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include "../../device.hpp" +#include "../../graph/graph.hpp" +#include "../common/op.hpp" + +#include + +namespace infinicore::op::distributed { + +class Broadcast : public graph::GraphOperator { +public: + Broadcast(Tensor output, const Tensor &input, int root, infinicclComm_t communicator); + ~Broadcast(); + void run() const override; + static void execute(Tensor output, const Tensor &input, int root, infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +Tensor broadcast(const Tensor &input, int root, infinicclComm_t communicator); +void broadcast_(Tensor output, const Tensor &input, int root, infinicclComm_t communicator); + +} // namespace infinicore::op::distributed diff --git a/include/infinicore/ops/distributed/send_recv.hpp b/include/infinicore/ops/distributed/send_recv.hpp new file mode 100644 index 000000000..c5c6c29d0 --- /dev/null +++ b/include/infinicore/ops/distributed/send_recv.hpp @@ -0,0 +1,37 @@ +#pragma once + +#include "../../device.hpp" +#include "../../graph/graph.hpp" +#include "../common/op.hpp" + +#include + +namespace infinicore::op::distributed { + +class Send : public graph::GraphOperator { +public: + Send(const Tensor &input, int peer, infinicclComm_t communicator); + ~Send(); + void run() const override; + static void execute(const Tensor &input, int peer, infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +class Recv : public graph::GraphOperator { +public: + Recv(Tensor output, int peer, infinicclComm_t communicator); + ~Recv(); + void run() const override; + static void execute(Tensor output, int peer, infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +void send(const Tensor &input, int peer, infinicclComm_t communicator); +void recv_(Tensor output, int peer, infinicclComm_t communicator); +Tensor recv(const Shape &shape, DataType dtype, Device device, int peer, infinicclComm_t communicator); + +} // namespace infinicore::op::distributed diff --git a/src/infiniccl-test/infiniccl_test.cpp b/src/infiniccl-test/infiniccl_test.cpp index 0aa898484..adc5c8a51 100644 --- a/src/infiniccl-test/infiniccl_test.cpp +++ b/src/infiniccl-test/infiniccl_test.cpp @@ -1,29 +1,38 @@ #include "infiniccl_test.hpp" +#include +#include +#include #include +#include +#include +#include #include #include +#include #include #include +#include +#include +#include +#include +#include +#include +#include #include #define TEST_INFINI(API__) CHECK_API_OR(API__, INFINI_STATUS_SUCCESS, return 1) #define TEST_INFINI_THREAD(API__) CHECK_API_OR(API__, INFINI_STATUS_SUCCESS, return nullptr) const size_t MAX_COUNT = 8ULL * 1024 * 1024; -// const size_t MAX_COUNT = 512 * 1024; // for metax - const size_t TEST_COUNTS[] = { 128, 1024, 4 * 1024, MAX_COUNT, }; - const infiniDtype_t TEST_DTYPES[] = {INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16}; - const size_t WARM_UPS = 10; - const size_t ITERATIONS = 100; struct ThreadArgs { @@ -39,6 +48,34 @@ struct ThreadArgs { double *time; }; +struct SendThreadArgs { + int rank; + int device_id; + int peer; + bool is_sender; + infinicclComm_t comm; + infiniDevice_t device_type; + infiniDtype_t dtype; + size_t count; + const void *data; + const void *ans; + int *result; + double *time; +}; + +struct BroadcastThreadArgs { + int rank; + int device_id; + infinicclComm_t comm; + infiniDevice_t device_type; + infiniDtype_t dtype; + size_t count; + const void *data; + const void *ans; + int *result; + double *time; +}; + void setData(infiniDtype_t dtype, void *data, size_t count, float val) { switch (dtype) { case INFINI_DTYPE_F32: @@ -46,7 +83,6 @@ void setData(infiniDtype_t dtype, void *data, size_t count, float val) { ((float *)data)[i] = val; } break; - case INFINI_DTYPE_F16: for (size_t i = 0; i < count; i++) { ((fp16_t *)data)[i] = utils::cast(val); @@ -67,20 +103,14 @@ template int checkData(const T *actual_, const T *expected_, size_t count) { int failed = 0; for (size_t i = 0; i < count; i++) { - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value || std::is_same::value) { float actual = utils::cast(actual_[i]); float expected = utils::cast(expected_[i]); - if (std::abs(actual - expected) > 1e-4) { - failed += 1; - } - } else if constexpr (std::is_same::value) { - float actual = utils::cast(actual_[i]); - float expected = utils::cast(expected_[i]); - if (std::abs(actual - expected) > 1e-4) { + if (std::abs(actual - expected) > 1e-4f) { failed += 1; } } else { - if (std::abs(actual_[i] - expected_[i]) > 1e-4) { + if (std::abs(actual_[i] - expected_[i]) > 1e-4f) { failed += 1; } } @@ -102,6 +132,13 @@ int checkData(const void *actual, const void *expected, infiniDtype_t dtype, siz } } +double bandwidthGBps(size_t bytes, double ms) { + if (ms <= 0.0) { + return 0.0; + } + return static_cast(bytes) / ms / 1.0e6; +} + void *testAllReduceThread(void *arg) { ThreadArgs *args = (ThreadArgs *)arg; *(args->result) = 1; @@ -128,7 +165,6 @@ void *testAllReduceThread(void *arg) { } TEST_INFINI_THREAD(infinirtStreamSynchronize(stream)); - // measure time auto start = std::chrono::high_resolution_clock::now(); for (size_t i = 0; i < ITERATIONS; i++) { TEST_INFINI_THREAD(infinicclAllReduce(buf, buf, args->count, args->dtype, INFINICCL_SUM, args->comm, stream)); @@ -137,7 +173,6 @@ void *testAllReduceThread(void *arg) { auto end = std::chrono::high_resolution_clock::now(); double elapsed_ms = std::chrono::duration(end - start).count(); *args->time = elapsed_ms / ITERATIONS; - *args->result = 0; std::free(output); @@ -153,7 +188,7 @@ int testAllReduce(infiniDevice_t device_type, int ndevice) { std::vector device_ids(ndevice); std::vector results(ndevice); std::vector times(ndevice); - void *data = std::malloc(MAX_COUNT * sizeof(float)); // Use float as max dtype size + void *data = std::malloc(MAX_COUNT * sizeof(float)); void *ans = std::malloc(MAX_COUNT * sizeof(float)); for (int i = 0; i < ndevice; i++) { @@ -178,8 +213,11 @@ int testAllReduce(infiniDevice_t device_type, int ndevice) { if (results[rank] != 0) { std::cout << "Rank " << rank << ": incorrect results." << std::endl; } else { - std::cout << "Rank " << rank << ": " << times[rank] << " ms." << std::endl; + auto bytes = count * infiniSizeOf(dtype); + std::cout << "Rank " << rank << ": " << times[rank] << " ms, " + << bandwidthGBps(bytes, times[rank]) << " GB/s payload." << std::endl; } + infinicclCommDestroy(comms[rank]); } if (failed > 0) { @@ -197,3 +235,438 @@ int testAllReduce(infiniDevice_t device_type, int ndevice) { std::free(ans); return 0; } + +void *testBroadcastThread(void *arg) { + BroadcastThreadArgs *args = (BroadcastThreadArgs *)arg; + *(args->result) = 1; + TEST_INFINI_THREAD(infinirtSetDevice(args->device_type, args->device_id)); + + infinirtStream_t stream; + TEST_INFINI_THREAD(infinirtStreamCreate(&stream)); + + const size_t bytes = args->count * infiniSizeOf(args->dtype); + void *buf = nullptr; + void *output = std::malloc(bytes); + TEST_INFINI_THREAD(infinirtMalloc(&buf, bytes)); + if (args->rank == 0) { + TEST_INFINI_THREAD(infinirtMemcpy(buf, args->data, bytes, INFINIRT_MEMCPY_H2D)); + } + + TEST_INFINI_THREAD(infinicclBroadcast(buf, buf, args->count, args->dtype, 0, args->comm, stream)); + TEST_INFINI_THREAD(infinirtStreamSynchronize(stream)); + TEST_INFINI_THREAD(infinirtMemcpy(output, buf, bytes, INFINIRT_MEMCPY_D2H)); + + int mismatches = checkData(output, args->ans, args->dtype, args->count); + if (mismatches != 0) { + std::cout << "Rank " << args->rank << ": broadcast correctness failed with " + << mismatches << " mismatches." << std::endl; + std::free(output); + infinirtFree(buf); + infinirtStreamDestroy(stream); + return nullptr; + } + + for (size_t i = 0; i < WARM_UPS; i++) { + TEST_INFINI_THREAD(infinicclBroadcast(buf, buf, args->count, args->dtype, 0, args->comm, stream)); + } + TEST_INFINI_THREAD(infinirtStreamSynchronize(stream)); + + auto start = std::chrono::high_resolution_clock::now(); + for (size_t i = 0; i < ITERATIONS; i++) { + TEST_INFINI_THREAD(infinicclBroadcast(buf, buf, args->count, args->dtype, 0, args->comm, stream)); + } + TEST_INFINI_THREAD(infinirtStreamSynchronize(stream)); + auto end = std::chrono::high_resolution_clock::now(); + double elapsed_ms = std::chrono::duration(end - start).count(); + *args->time = elapsed_ms / ITERATIONS; + *args->result = 0; + + std::free(output); + infinirtFree(buf); + infinirtStreamDestroy(stream); + return nullptr; +} + +int testBroadcast(infiniDevice_t device_type, int ndevice) { + std::vector thread_args(ndevice); + std::vector comms(ndevice); + std::vector threads(ndevice); + std::vector device_ids(ndevice); + std::vector results(ndevice); + std::vector times(ndevice); + void *data = std::malloc(MAX_COUNT * sizeof(float)); + void *ans = std::malloc(MAX_COUNT * sizeof(float)); + + for (int i = 0; i < ndevice; i++) { + device_ids[i] = i; + } + + for (infiniDtype_t dtype : TEST_DTYPES) { + setData(dtype, data, MAX_COUNT, 9.0f); + setData(dtype, ans, MAX_COUNT, 9.0f); + for (size_t count : TEST_COUNTS) { + TEST_INFINI(infinicclCommInitAll(device_type, comms.data(), ndevice, device_ids.data())); + std::cout << "Testing Broadcast root rank 0 with " << count << " elements of " + << infiniDtypeToString(dtype) << std::endl; + + for (int rank = 0; rank < ndevice; rank++) { + thread_args[rank] = {rank, device_ids[rank], comms[rank], device_type, dtype, count, + data, ans, &results[rank], ×[rank]}; + pthread_create(&threads[rank], NULL, testBroadcastThread, &thread_args[rank]); + } + for (int rank = 0; rank < ndevice; rank++) { + pthread_join(threads[rank], NULL); + } + + int failed = std::accumulate(results.begin(), results.end(), 0); + for (int rank = 0; rank < ndevice; rank++) { + if (results[rank] != 0) { + std::cout << "Rank " << rank << ": failed." << std::endl; + } else { + auto bytes = count * infiniSizeOf(dtype); + std::cout << "Rank " << rank << ": " << times[rank] << " ms, " + << bandwidthGBps(bytes, times[rank]) << " GB/s payload." << std::endl; + } + infinicclCommDestroy(comms[rank]); + } + + if (failed > 0) { + std::free(data); + std::free(ans); + return 1; + } + std::cout << std::endl; + } + } + + std::free(data); + std::free(ans); + return 0; +} + +void *testSendThread(void *arg) { + SendThreadArgs *args = (SendThreadArgs *)arg; + *(args->result) = 1; + TEST_INFINI_THREAD(infinirtSetDevice(args->device_type, args->device_id)); + + infinirtStream_t stream; + TEST_INFINI_THREAD(infinirtStreamCreate(&stream)); + + const size_t bytes = args->count * infiniSizeOf(args->dtype); + void *buf = nullptr; + void *output = std::malloc(bytes); + TEST_INFINI_THREAD(infinirtMalloc(&buf, bytes)); + if (args->is_sender) { + TEST_INFINI_THREAD(infinirtMemcpy(buf, args->data, bytes, INFINIRT_MEMCPY_H2D)); + } + + if (args->is_sender) { + TEST_INFINI_THREAD(infinicclSend(buf, args->count, args->dtype, args->peer, args->comm, stream)); + } else { + TEST_INFINI_THREAD(infinicclRecv(buf, args->count, args->dtype, args->peer, args->comm, stream)); + } + TEST_INFINI_THREAD(infinirtStreamSynchronize(stream)); + + int mismatches = 0; + if (!args->is_sender) { + TEST_INFINI_THREAD(infinirtMemcpy(output, buf, bytes, INFINIRT_MEMCPY_D2H)); + mismatches = checkData(output, args->ans, args->dtype, args->count); + if (mismatches != 0) { + std::cout << "Rank " << args->rank << ": send/recv correctness failed with " + << mismatches << " mismatches." << std::endl; + } + } + + for (size_t i = 0; i < WARM_UPS; i++) { + if (args->is_sender) { + TEST_INFINI_THREAD(infinicclSend(buf, args->count, args->dtype, args->peer, args->comm, stream)); + } else { + TEST_INFINI_THREAD(infinicclRecv(buf, args->count, args->dtype, args->peer, args->comm, stream)); + } + } + TEST_INFINI_THREAD(infinirtStreamSynchronize(stream)); + + auto start = std::chrono::high_resolution_clock::now(); + for (size_t i = 0; i < ITERATIONS; i++) { + if (args->is_sender) { + TEST_INFINI_THREAD(infinicclSend(buf, args->count, args->dtype, args->peer, args->comm, stream)); + } else { + TEST_INFINI_THREAD(infinicclRecv(buf, args->count, args->dtype, args->peer, args->comm, stream)); + } + } + TEST_INFINI_THREAD(infinirtStreamSynchronize(stream)); + auto end = std::chrono::high_resolution_clock::now(); + double elapsed_ms = std::chrono::duration(end - start).count(); + *args->time = elapsed_ms / ITERATIONS; + *args->result = mismatches == 0 ? 0 : 1; + + std::free(output); + infinirtFree(buf); + infinirtStreamDestroy(stream); + return nullptr; +} + +int testSend(infiniDevice_t device_type, int ndevice, bool send_from_zero) { + if (ndevice < 2) { + std::cout << "Send/recv test requires at least 2 visible devices." << std::endl; + return 1; + } + + std::vector device_ids = {0, 1}; + std::vector comms(2); + std::vector threads(2); + std::vector results(2); + std::vector times(2); + std::vector thread_args(2); + const int src_rank = send_from_zero ? 0 : 1; + const int dst_rank = send_from_zero ? 1 : 0; + + for (infiniDtype_t dtype : TEST_DTYPES) { + for (size_t count : TEST_COUNTS) { + const size_t bytes = count * infiniSizeOf(dtype); + void *data = std::malloc(bytes); + void *ans = std::malloc(bytes); + setData(dtype, data, count, 7.0f); + setData(dtype, ans, count, 7.0f); + + TEST_INFINI(infinicclCommInitAll(device_type, comms.data(), 2, device_ids.data())); + std::cout << "Testing Send rank " << src_rank << " -> rank " << dst_rank + << " with " << count << " elements of " << infiniDtypeToString(dtype) << std::endl; + + thread_args[0] = {0, device_ids[0], 1, send_from_zero, comms[0], device_type, dtype, count, data, ans, &results[0], ×[0]}; + thread_args[1] = {1, device_ids[1], 0, !send_from_zero, comms[1], device_type, dtype, count, data, ans, &results[1], ×[1]}; + pthread_create(&threads[0], NULL, testSendThread, &thread_args[0]); + pthread_create(&threads[1], NULL, testSendThread, &thread_args[1]); + pthread_join(threads[0], NULL); + pthread_join(threads[1], NULL); + + int failed = results[0] + results[1]; + for (int rank = 0; rank < 2; rank++) { + if (results[rank] != 0) { + std::cout << "Rank " << rank << ": failed." << std::endl; + } else { + std::cout << "Rank " << rank << ": " << times[rank] << " ms, " + << bandwidthGBps(bytes, times[rank]) << " GB/s payload." << std::endl; + } + infinicclCommDestroy(comms[rank]); + } + + std::free(data); + std::free(ans); + if (failed > 0) { + return 1; + } + std::cout << std::endl; + } + } + return 0; +} + +namespace { + +void throwErrno(const char *what) { + throw std::runtime_error(std::string(what) + ": " + std::strerror(errno)); +} + +void sendAll(int fd, const void *data, size_t size) { + const char *ptr = static_cast(data); + while (size > 0) { + ssize_t n = send(fd, ptr, size, 0); + if (n <= 0) { + throwErrno("send"); + } + ptr += n; + size -= static_cast(n); + } +} + +void recvAll(int fd, void *data, size_t size) { + char *ptr = static_cast(data); + while (size > 0) { + ssize_t n = recv(fd, ptr, size, 0); + if (n <= 0) { + throwErrno("recv"); + } + ptr += n; + size -= static_cast(n); + } +} + +void exchangeUniqueId( + int rank, + int world_size, + const char *master_addr, + int master_port, + infinicclUniqueId_t *unique_id) { + if (rank == 0) { + int listen_fd = socket(AF_INET, SOCK_STREAM, 0); + if (listen_fd < 0) { + throwErrno("socket"); + } + int yes = 1; + setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes)); + + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_addr.s_addr = htonl(INADDR_ANY); + addr.sin_port = htons(static_cast(master_port)); + if (bind(listen_fd, reinterpret_cast(&addr), sizeof(addr)) != 0) { + close(listen_fd); + throwErrno("bind"); + } + if (listen(listen_fd, world_size - 1) != 0) { + close(listen_fd); + throwErrno("listen"); + } + + for (int i = 1; i < world_size; ++i) { + int client_fd = accept(listen_fd, nullptr, nullptr); + if (client_fd < 0) { + close(listen_fd); + throwErrno("accept"); + } + sendAll(client_fd, unique_id, sizeof(*unique_id)); + close(client_fd); + } + close(listen_fd); + } else { + int fd = -1; + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = htons(static_cast(master_port)); + if (inet_pton(AF_INET, master_addr, &addr.sin_addr) != 1) { + throw std::runtime_error("master-addr must be an IPv4 address"); + } + + for (int attempt = 0; attempt < 300; ++attempt) { + fd = socket(AF_INET, SOCK_STREAM, 0); + if (fd < 0) { + throwErrno("socket"); + } + if (connect(fd, reinterpret_cast(&addr), sizeof(addr)) == 0) { + break; + } + close(fd); + fd = -1; + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + if (fd < 0) { + throw std::runtime_error("failed to connect to rank 0 unique-id server"); + } + recvAll(fd, unique_id, sizeof(*unique_id)); + close(fd); + } +} + +} // namespace + +int testMultiNodeSend( + infiniDevice_t device_type, + bool is_sender, + const char *master_addr, + int master_port) { + try { + const int world_size = 2; + const bool is_master = master_addr == nullptr; + const int rank = is_master ? 0 : 1; + const int peer = 1 - rank; + const int device_id = 0; + const char *connect_addr = is_master ? "0.0.0.0" : master_addr; + + TEST_INFINI(infinirtSetDevice(device_type, device_id)); + + infinicclUniqueId_t unique_id; + if (rank == 0) { + TEST_INFINI(infinicclGetUniqueId(&unique_id)); + } + + std::cout << "Global rank " << rank << "/" << world_size + << " exchanging unique id via " << connect_addr << ":" << master_port << std::endl; + exchangeUniqueId(rank, world_size, connect_addr, master_port, &unique_id); + + infinicclComm_t comm = nullptr; + TEST_INFINI(infinicclCommInitRank(&comm, world_size, unique_id, rank)); + + infinirtStream_t stream; + TEST_INFINI(infinirtStreamCreate(&stream)); + + int failed = 0; + for (infiniDtype_t dtype : TEST_DTYPES) { + for (size_t count : TEST_COUNTS) { + const size_t bytes = count * infiniSizeOf(dtype); + void *data = std::malloc(bytes); + void *ans = std::malloc(bytes); + void *output = std::malloc(bytes); + void *buf = nullptr; + setData(dtype, data, count, 7.0f); + setData(dtype, ans, count, 7.0f); + std::memset(output, 0, bytes); + + TEST_INFINI(infinirtMalloc(&buf, bytes)); + if (is_sender) { + TEST_INFINI(infinirtMemcpy(buf, data, bytes, INFINIRT_MEMCPY_H2D)); + } + + if (is_sender) { + TEST_INFINI(infinicclSend(buf, count, dtype, peer, comm, stream)); + } else { + TEST_INFINI(infinicclRecv(buf, count, dtype, peer, comm, stream)); + } + TEST_INFINI(infinirtStreamSynchronize(stream)); + + int mismatches = 0; + if (!is_sender) { + TEST_INFINI(infinirtMemcpy(output, buf, bytes, INFINIRT_MEMCPY_D2H)); + mismatches = checkData(output, ans, dtype, count); + if (mismatches != 0) { + std::cout << "Global rank " << rank << " " << infiniDtypeToString(dtype) + << " count=" << count << " failed with " << mismatches + << " mismatches." << std::endl; + } + } + + for (size_t i = 0; i < WARM_UPS; i++) { + if (is_sender) { + TEST_INFINI(infinicclSend(buf, count, dtype, peer, comm, stream)); + } else { + TEST_INFINI(infinicclRecv(buf, count, dtype, peer, comm, stream)); + } + } + TEST_INFINI(infinirtStreamSynchronize(stream)); + + auto start = std::chrono::high_resolution_clock::now(); + for (size_t i = 0; i < ITERATIONS; i++) { + if (is_sender) { + TEST_INFINI(infinicclSend(buf, count, dtype, peer, comm, stream)); + } else { + TEST_INFINI(infinicclRecv(buf, count, dtype, peer, comm, stream)); + } + } + TEST_INFINI(infinirtStreamSynchronize(stream)); + auto end = std::chrono::high_resolution_clock::now(); + double elapsed_ms = std::chrono::duration(end - start).count(); + double avg_ms = elapsed_ms / ITERATIONS; + + infinirtFree(buf); + std::free(data); + std::free(ans); + std::free(output); + + if (mismatches != 0) { + failed = 1; + } + std::cout << "Global rank " << rank << " " << (is_sender ? "send" : "recv") + << " " << infiniDtypeToString(dtype) << " count=" << count + << ", avg=" << avg_ms << " ms, " << bandwidthGBps(bytes, avg_ms) + << " GB/s payload." << std::endl; + } + } + + infinirtStreamDestroy(stream); + infinicclCommDestroy(comm); + return failed; + } catch (const std::exception &e) { + std::cout << "Global rank " << (is_sender ? 0 : 1) << " failed: " << e.what() << std::endl; + return 1; + } +} diff --git a/src/infiniccl-test/infiniccl_test.hpp b/src/infiniccl-test/infiniccl_test.hpp index 0ea4bbed6..848becad4 100644 --- a/src/infiniccl-test/infiniccl_test.hpp +++ b/src/infiniccl-test/infiniccl_test.hpp @@ -6,5 +6,12 @@ #include "../utils.h" int testAllReduce(infiniDevice_t device_type, int ndevice); +int testBroadcast(infiniDevice_t device_type, int ndevice); +int testSend(infiniDevice_t device_type, int ndevice, bool send_from_zero); +int testMultiNodeSend( + infiniDevice_t device_type, + bool is_sender, + const char *master_addr, + int master_port); #endif // INFINICCL_TEST_HPP diff --git a/src/infiniccl-test/main.cpp b/src/infiniccl-test/main.cpp index 8126aa15a..df97d0480 100644 --- a/src/infiniccl-test/main.cpp +++ b/src/infiniccl-test/main.cpp @@ -1,41 +1,64 @@ #include "infiniccl_test.hpp" +#include #include +#include + +enum class TestFunc { + AllReduce, + Broadcast, + Send, + Recv, +}; struct ParsedArgs { - infiniDevice_t device_type; + infiniDevice_t device_type = INFINI_DEVICE_CPU; + bool parsed_device = false; + bool multi_node = false; + TestFunc func = TestFunc::AllReduce; + bool has_master_addr = false; + std::string master_addr; + int master_port = 29500; }; void printUsage() { std::cout << "Usage:" << std::endl << std::endl; - std::cout << "infiniccl-test --" << std::endl + std::cout << "infiniccl-test -- [--func allreduce|broadcast|send|recv]" << std::endl + << "infiniccl-test -- --multi-node --func send|recv " + << "[--master-addr ] [--master-port ]" << std::endl << std::endl; std::cout << " --" << std::endl; std::cout << " Specify the device type --(nvidia|cambricon|ascend|metax|moore|iluvatar|qy|kunlun|hygon|ali)." << std::endl << std::endl; - std::cout << "The program will run tests on all visible devices of the specified device type." - << " Use Environmental Variables such as CUDA_VSIBLE_DEVICES to limit visible device IDs."; + std::cout << " --func" << std::endl + << " allreduce: local all-reduce test across all visible devices." << std::endl + << " broadcast: local broadcast test from device 0 to all visible devices." << std::endl + << " send: local device 0 sends to device 1, or multi-node sender on visible device 0." << std::endl + << " recv: local device 1 sends to device 0, or multi-node receiver on visible device 0." << std::endl + << std::endl; + std::cout << "For multi-node send/recv, the process without --master-addr is the TCP master." + << std::endl; exit(-1); } #define PARSE_DEVICE(FLAG, DEVICE) \ if (arg == FLAG) { \ args.device_type = DEVICE; \ + args.parsed_device = true; \ } ParsedArgs parseArgs(int argc, char *argv[]) { - if (argc != 2) { - printUsage(); - } - - if (std::string(argv[1]) == "--help" || std::string(argv[1]) == "-h") { + if (argc < 2) { printUsage(); } ParsedArgs args; - try { - std::string arg = argv[1]; + for (int i = 1; i < argc; ++i) { + std::string arg = argv[i]; + if (arg == "--help" || arg == "-h") { + printUsage(); + } // clang-format off PARSE_DEVICE("--nvidia", INFINI_DEVICE_NVIDIA) else PARSE_DEVICE("--cambricon", INFINI_DEVICE_CAMBRICON) @@ -47,22 +70,55 @@ ParsedArgs parseArgs(int argc, char *argv[]) { else PARSE_DEVICE("--kunlun", INFINI_DEVICE_KUNLUN) else PARSE_DEVICE("--hygon", INFINI_DEVICE_HYGON) else PARSE_DEVICE("--ali", INFINI_DEVICE_ALI) - else { + else if (arg == "--multi-node") { + args.multi_node = true; + } else if (arg == "--func" && i + 1 < argc) { + std::string func = argv[++i]; + if (func == "allreduce") { + args.func = TestFunc::AllReduce; + } else if (func == "broadcast") { + args.func = TestFunc::Broadcast; + } else if (func == "send") { + args.func = TestFunc::Send; + } else if (func == "recv") { + args.func = TestFunc::Recv; + } else { + printUsage(); + } + } else if (arg == "--master-addr" && i + 1 < argc) { + args.has_master_addr = true; + args.master_addr = argv[++i]; + } else if (arg == "--master-port" && i + 1 < argc) { + args.master_port = std::atoi(argv[++i]); + } else { printUsage(); } // clang-format on + } - } catch (const std::exception &) { + if (!args.parsed_device || args.master_port <= 0 || args.master_port > 65535) { + printUsage(); + } + if (args.multi_node && (args.func == TestFunc::AllReduce || args.func == TestFunc::Broadcast)) { + printUsage(); + } + if (!args.multi_node && args.has_master_addr) { printUsage(); } - return args; } int main(int argc, char *argv[]) { ParsedArgs args = parseArgs(argc, argv); - int ndevice = 0; infinirtInit(); + + if (args.multi_node) { + return testMultiNodeSend(args.device_type, args.func == TestFunc::Send, + args.has_master_addr ? args.master_addr.c_str() : nullptr, + args.master_port); + } + + int ndevice = 0; if (infinirtGetDeviceCount(args.device_type, &ndevice) != INFINI_STATUS_SUCCESS) { std::cout << "Failed to get device count" << std::endl; return -1; @@ -70,11 +126,17 @@ int main(int argc, char *argv[]) { if (ndevice == 0) { std::cout << "No devices found. Tests skipped." << std::endl; return 0; - } else { - std::cout << "Found " << ndevice << " devices. Running tests..." << std::endl; } - int failed = 0; - failed += testAllReduce(args.device_type, ndevice); - return failed; + std::cout << "Found " << ndevice << " devices. Running tests..." << std::endl; + if (args.func == TestFunc::Send) { + return testSend(args.device_type, ndevice, true); + } + if (args.func == TestFunc::Recv) { + return testSend(args.device_type, ndevice, false); + } + if (args.func == TestFunc::Broadcast) { + return testBroadcast(args.device_type, ndevice); + } + return testAllReduce(args.device_type, ndevice); } diff --git a/src/infiniccl/cuda/infiniccl_cuda.cu b/src/infiniccl/cuda/infiniccl_cuda.cu index 6a8442d21..30f68f636 100644 --- a/src/infiniccl/cuda/infiniccl_cuda.cu +++ b/src/infiniccl/cuda/infiniccl_cuda.cu @@ -1,5 +1,6 @@ #include "infiniccl_cuda.h" +#include #include #include #include @@ -77,6 +78,39 @@ infiniStatus_t commInitAll( return INFINI_STATUS_SUCCESS; } +infiniStatus_t getUniqueId(infinicclUniqueId_t *unique_id) { + if (unique_id == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + + CHECK_NCCL(ncclGetUniqueId(reinterpret_cast(unique_id))); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t commInitRank( + infinicclComm_t *comm, + int nranks, + infinicclUniqueId_t comm_id, + int rank) { + if (comm == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + if (nranks <= 0 || rank < 0 || rank >= nranks) { + return INFINI_STATUS_BAD_PARAM; + } + + infiniDevice_t device_type; + int device_id; + CHECK_STATUS(infinirtGetDevice(&device_type, &device_id)); + + ncclUniqueId nccl_id; + std::memcpy(&nccl_id, &comm_id, sizeof(nccl_id)); + ncclComm_t nccl_comm; + CHECK_NCCL(ncclCommInitRank(&nccl_comm, nranks, nccl_id, rank)); + *comm = new InfinicclComm{device_type, device_id, (void *)nccl_comm, rank, nranks}; + return INFINI_STATUS_SUCCESS; +} + infiniStatus_t commDestroy(infinicclComm_t comm) { CHECK_NCCL(ncclCommDestroy(getNcclComm(comm))); delete comm; @@ -110,6 +144,76 @@ infiniStatus_t allReduce( return INFINI_STATUS_SUCCESS; } +infiniStatus_t broadcast( + const void *sendbuf, + void *recvbuf, + size_t count, + infiniDtype_t datatype, + int root, + infinicclComm_t comm, + infinirtStream_t stream) { + + if (sendbuf == nullptr || recvbuf == nullptr || comm == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + if (root < 0 || root >= comm->world_size || count == 0) { + return INFINI_STATUS_BAD_PARAM; + } + CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, + INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_U32, INFINI_DTYPE_U64); + + CHECK_NCCL(ncclBroadcast(sendbuf, recvbuf, count, getNcclDtype(datatype), root, + getNcclComm(comm), getCudaStream(stream))); + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t send( + const void *sendbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream) { + + if (sendbuf == nullptr || comm == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + if (peer < 0 || peer >= comm->world_size || count == 0) { + return INFINI_STATUS_BAD_PARAM; + } + CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, + INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_U32, INFINI_DTYPE_U64); + + CHECK_NCCL(ncclSend(sendbuf, count, getNcclDtype(datatype), peer, + getNcclComm(comm), getCudaStream(stream))); + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t recv( + void *recvbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream) { + + if (recvbuf == nullptr || comm == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + if (peer < 0 || peer >= comm->world_size || count == 0) { + return INFINI_STATUS_BAD_PARAM; + } + CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, + INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_U32, INFINI_DTYPE_U64); + + CHECK_NCCL(ncclRecv(recvbuf, count, getNcclDtype(datatype), peer, + getNcclComm(comm), getCudaStream(stream))); + + return INFINI_STATUS_SUCCESS; +} + infiniStatus_t allGather( void *sendbuf, void *recvbuf, diff --git a/src/infiniccl/cuda/infiniccl_cuda.h b/src/infiniccl/cuda/infiniccl_cuda.h index 9bba02345..d8bc05105 100644 --- a/src/infiniccl/cuda/infiniccl_cuda.h +++ b/src/infiniccl/cuda/infiniccl_cuda.h @@ -6,8 +6,78 @@ // Windows does not support CUDA #if (defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API)) && defined(ENABLE_CCL) && !defined(_WIN32) INFINICCL_DEVICE_API_IMPL(cuda) +namespace infiniccl::cuda { +infiniStatus_t getUniqueId(infinicclUniqueId_t *unique_id); +infiniStatus_t commInitRank( + infinicclComm_t *comm, + int nranks, + infinicclUniqueId_t comm_id, + int rank); +infiniStatus_t broadcast( + const void *sendbuf, + void *recvbuf, + size_t count, + infiniDtype_t datatype, + int root, + infinicclComm_t comm, + infinirtStream_t stream); +infiniStatus_t send( + const void *sendbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream); +infiniStatus_t recv( + void *recvbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream); +} // namespace infiniccl::cuda #else INFINICCL_DEVICE_API_NOOP(cuda) +namespace infiniccl::cuda { +inline infiniStatus_t getUniqueId(infinicclUniqueId_t *) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} +inline infiniStatus_t commInitRank( + infinicclComm_t *, + int, + infinicclUniqueId_t, + int) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} +inline infiniStatus_t broadcast( + const void *, + void *, + size_t, + infiniDtype_t, + int, + infinicclComm_t, + infinirtStream_t) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} +inline infiniStatus_t send( + const void *, + size_t, + infiniDtype_t, + int, + infinicclComm_t, + infinirtStream_t) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} +inline infiniStatus_t recv( + void *, + size_t, + infiniDtype_t, + int, + infinicclComm_t, + infinirtStream_t) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} +} // namespace infiniccl::cuda #endif #endif /* INFINICCL_CUDA_H_ */ diff --git a/src/infiniccl/infiniccl.cc b/src/infiniccl/infiniccl.cc index 48dbea425..f348e355f 100644 --- a/src/infiniccl/infiniccl.cc +++ b/src/infiniccl/infiniccl.cc @@ -1,5 +1,6 @@ #include "infiniccl.h" +#include "../utils.h" #include "./ascend/infiniccl_ascend.h" #include "./cambricon/infiniccl_cambricon.h" #include "./cuda/infiniccl_cuda.h" @@ -35,6 +36,61 @@ __INFINI_C infiniStatus_t infinicclCommInitAll( #undef COMM_INIT_ALL } +__INFINI_C infiniStatus_t infinicclGetUniqueId( + infinicclUniqueId_t *unique_id) { + if (unique_id == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + + infiniDevice_t device_type; + CHECK_STATUS(infinirtGetDevice(&device_type, nullptr)); + +#define GET_UNIQUE_ID(CASE_, NAMESPACE_) \ + case CASE_: \ + return infiniccl::NAMESPACE_::getUniqueId(unique_id) + + switch (device_type) { + GET_UNIQUE_ID(INFINI_DEVICE_NVIDIA, cuda); + GET_UNIQUE_ID(INFINI_DEVICE_ILUVATAR, cuda); + GET_UNIQUE_ID(INFINI_DEVICE_QY, cuda); + GET_UNIQUE_ID(INFINI_DEVICE_HYGON, cuda); + GET_UNIQUE_ID(INFINI_DEVICE_ALI, cuda); + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef GET_UNIQUE_ID +} + +__INFINI_C infiniStatus_t infinicclCommInitRank( + infinicclComm_t *comm, + int nranks, + infinicclUniqueId_t comm_id, + int rank) { + if (comm == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + + infiniDevice_t device_type; + CHECK_STATUS(infinirtGetDevice(&device_type, nullptr)); + +#define COMM_INIT_RANK(CASE_, NAMESPACE_) \ + case CASE_: \ + return infiniccl::NAMESPACE_::commInitRank(comm, nranks, comm_id, rank) + + switch (device_type) { + COMM_INIT_RANK(INFINI_DEVICE_NVIDIA, cuda); + COMM_INIT_RANK(INFINI_DEVICE_ILUVATAR, cuda); + COMM_INIT_RANK(INFINI_DEVICE_QY, cuda); + COMM_INIT_RANK(INFINI_DEVICE_HYGON, cuda); + COMM_INIT_RANK(INFINI_DEVICE_ALI, cuda); + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef COMM_INIT_RANK +} + __INFINI_C infiniStatus_t infinicclCommDestroy(infinicclComm_t comm) { if (comm == nullptr) { return INFINI_STATUS_SUCCESS; @@ -151,6 +207,94 @@ __INFINI_C infiniStatus_t infinicclAllReduce( #undef ALL_REDUCE } +__INFINI_C infiniStatus_t infinicclBroadcast( + const void *sendbuf, + void *recvbuf, + size_t count, + infiniDtype_t datatype, + int root, + infinicclComm_t comm, + infinirtStream_t stream) { + + if (comm == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + +#define BROADCAST(CASE_, NAMESPACE_) \ + case CASE_: \ + return infiniccl::NAMESPACE_::broadcast(sendbuf, recvbuf, count, datatype, root, comm, stream) + + switch (comm->device_type) { + BROADCAST(INFINI_DEVICE_NVIDIA, cuda); + BROADCAST(INFINI_DEVICE_ILUVATAR, cuda); + BROADCAST(INFINI_DEVICE_QY, cuda); + BROADCAST(INFINI_DEVICE_HYGON, cuda); + BROADCAST(INFINI_DEVICE_ALI, cuda); + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef BROADCAST +} + +__INFINI_C infiniStatus_t infinicclSend( + const void *sendbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream) { + + if (comm == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + +#define SEND(CASE_, NAMESPACE_) \ + case CASE_: \ + return infiniccl::NAMESPACE_::send(sendbuf, count, datatype, peer, comm, stream) + + switch (comm->device_type) { + SEND(INFINI_DEVICE_NVIDIA, cuda); + SEND(INFINI_DEVICE_ILUVATAR, cuda); + SEND(INFINI_DEVICE_QY, cuda); + SEND(INFINI_DEVICE_HYGON, cuda); + SEND(INFINI_DEVICE_ALI, cuda); + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef SEND +} + +__INFINI_C infiniStatus_t infinicclRecv( + void *recvbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream) { + + if (comm == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + +#define RECV(CASE_, NAMESPACE_) \ + case CASE_: \ + return infiniccl::NAMESPACE_::recv(recvbuf, count, datatype, peer, comm, stream) + + switch (comm->device_type) { + RECV(INFINI_DEVICE_NVIDIA, cuda); + RECV(INFINI_DEVICE_ILUVATAR, cuda); + RECV(INFINI_DEVICE_QY, cuda); + RECV(INFINI_DEVICE_HYGON, cuda); + RECV(INFINI_DEVICE_ALI, cuda); + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef RECV +} + __INFINI_C infiniStatus_t infinicclAllGather( void *sendbuf, void *recvbuf, diff --git a/src/infinicore/ops/distributed/broadcast.cc b/src/infinicore/ops/distributed/broadcast.cc new file mode 100644 index 000000000..f44c63328 --- /dev/null +++ b/src/infinicore/ops/distributed/broadcast.cc @@ -0,0 +1,54 @@ +#include "infinicore/ops/distributed/broadcast.hpp" +#include "../../utils.hpp" + +#include "infinicore/context/context.hpp" + +namespace infinicore::op::distributed { + +struct BroadcastPlannedMeta { + graph::GraphTensor output, input; + int root; + infinicclComm_t communicator; +}; + +Broadcast::Broadcast(Tensor output, const Tensor &input, int root, infinicclComm_t communicator) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); + INFINICORE_ASSERT(output->numel() == input->numel()); + INFINICORE_ASSERT(input->numel() > 0); + planned_meta_ = new BroadcastPlannedMeta{graph::GraphTensor(output), graph::GraphTensor(input), root, communicator}; +} + +Broadcast::~Broadcast() { + if (planned_meta_) { + auto *meta = reinterpret_cast(planned_meta_); + delete meta; + } +} + +void Broadcast::run() const { + auto *meta = reinterpret_cast(planned_meta_); + INFINICORE_CHECK_ERROR(infinicclBroadcast(meta->input->data(), + meta->output->data(), + meta->input->numel(), + static_cast(static_cast(meta->input->dtype())), + meta->root, + meta->communicator, + infinicore::context::getStream())); +} + +void Broadcast::execute(Tensor output, const Tensor &input, int root, infinicclComm_t communicator) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Broadcast, output, input, root, communicator); +} + +Tensor broadcast(const Tensor &input, int root, infinicclComm_t communicator) { + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + broadcast_(output, input, root, communicator); + return output; +} + +void broadcast_(Tensor output, const Tensor &input, int root, infinicclComm_t communicator) { + Broadcast::execute(output, input, root, communicator); +} + +} // namespace infinicore::op::distributed diff --git a/src/infinicore/ops/distributed/send_recv.cc b/src/infinicore/ops/distributed/send_recv.cc new file mode 100644 index 000000000..8a7998ed3 --- /dev/null +++ b/src/infinicore/ops/distributed/send_recv.cc @@ -0,0 +1,88 @@ +#include "infinicore/ops/distributed/send_recv.hpp" +#include "../../utils.hpp" + +#include "infinicore/context/context.hpp" + +namespace infinicore::op::distributed { + +struct SendPlannedMeta { + graph::GraphTensor input; + int peer; + infinicclComm_t communicator; +}; + +struct RecvPlannedMeta { + graph::GraphTensor output; + int peer; + infinicclComm_t communicator; +}; + +Send::Send(const Tensor &input, int peer, infinicclComm_t communicator) { + INFINICORE_ASSERT(input->is_contiguous()); + INFINICORE_ASSERT(input->numel() > 0); + planned_meta_ = new SendPlannedMeta{graph::GraphTensor(input), peer, communicator}; +} + +Send::~Send() { + if (planned_meta_) { + auto *meta = reinterpret_cast(planned_meta_); + delete meta; + } +} + +void Send::run() const { + auto *meta = reinterpret_cast(planned_meta_); + INFINICORE_CHECK_ERROR(infinicclSend(meta->input->data(), + meta->input->numel(), + static_cast(static_cast(meta->input->dtype())), + meta->peer, + meta->communicator, + infinicore::context::getStream())); +} + +void Send::execute(const Tensor &input, int peer, infinicclComm_t communicator) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Send, input, peer, communicator); +} + +Recv::Recv(Tensor output, int peer, infinicclComm_t communicator) { + INFINICORE_ASSERT(output->is_contiguous()); + INFINICORE_ASSERT(output->numel() > 0); + planned_meta_ = new RecvPlannedMeta{graph::GraphTensor(output), peer, communicator}; +} + +Recv::~Recv() { + if (planned_meta_) { + auto *meta = reinterpret_cast(planned_meta_); + delete meta; + } +} + +void Recv::run() const { + auto *meta = reinterpret_cast(planned_meta_); + INFINICORE_CHECK_ERROR(infinicclRecv(meta->output->data(), + meta->output->numel(), + static_cast(static_cast(meta->output->dtype())), + meta->peer, + meta->communicator, + infinicore::context::getStream())); +} + +void Recv::execute(Tensor output, int peer, infinicclComm_t communicator) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Recv, output, peer, communicator); +} + +void send(const Tensor &input, int peer, infinicclComm_t communicator) { + Send::execute(input, peer, communicator); +} + +void recv_(Tensor output, int peer, infinicclComm_t communicator) { + Recv::execute(output, peer, communicator); +} + +Tensor recv(const Shape &shape, DataType dtype, Device device, int peer, infinicclComm_t communicator) { + auto output = Tensor::empty(shape, dtype, device); + recv_(output, peer, communicator); + return output; +} + +} // namespace infinicore::op::distributed