feat InfiniCCL initRank/Send/Recv/Broadcast

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,

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"

include/infinicore/ops/distributed/broadcast.hpp

@@ -0,0 +1,25 @@
+#pragma once
+
+#include "../../device.hpp"
+#include "../../graph/graph.hpp"
+#include "../common/op.hpp"
+
+#include <infiniccl.h>
+
+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

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 <infiniccl.h>
+
+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