Merge pull request #257 from InfiniTensor/issue/256

issue/256 沐曦通信库

src/infiniccl/infiniccl.cc

@@ -2,6 +2,7 @@
 
 #include "./ascend/infiniccl_ascend.h"
 #include "./cuda/infiniccl_cuda.h"
+#include "./maca/infiniccl_maca.h"
 
 __C infiniStatus_t infinicclCommInitAll(
     infiniDevice_t device_type,
@@ -16,6 +17,7 @@ __C infiniStatus_t infinicclCommInitAll(
     switch (device_type) {
         COMM_INIT_ALL(INFINI_DEVICE_NVIDIA, cuda)
         COMM_INIT_ALL(INFINI_DEVICE_ASCEND, ascend)
+        COMM_INIT_ALL(INFINI_DEVICE_METAX, maca)
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
     }
@@ -35,6 +37,8 @@ __C infiniStatus_t infinicclCommDestroy(infinicclComm_t comm) {
     switch (comm->device_type) {
         COMM_DESTROY(INFINI_DEVICE_NVIDIA, cuda)
         COMM_DESTROY(INFINI_DEVICE_ASCEND, ascend)
+        COMM_DESTROY(INFINI_DEVICE_METAX, maca)
+
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
     }
@@ -61,6 +65,8 @@ __C infiniStatus_t infinicclAllReduce(
     switch (comm->device_type) {
         ALL_REDUCE(INFINI_DEVICE_NVIDIA, cuda)
         ALL_REDUCE(INFINI_DEVICE_ASCEND, ascend)
+        ALL_REDUCE(INFINI_DEVICE_METAX, maca)
+
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
     }

src/infiniccl/maca/infiniccl_maca.cc

+#include "infiniccl_maca.h"
+
+#include "../../utils.h"
+
+#include <hccl.h>
+#include <hcr/hc_runtime_api.h>
+
+#include <iostream>
+#include <vector>
+
+#define CHECK_HCCL(API__) CHECK_INTERNAL(API__, hcclSuccess)
+
+inline hcStream_t getMacaStream(infinirtStream_t stream) {
+    if (stream == nullptr) {
+        return 0;
+    }
+    return static_cast<hcStream_t>(stream);
+}
+
+inline hcclDataType_t getHcclDtype(infiniDtype_t datatype) {
+    switch (datatype) {
+    case INFINI_DTYPE_F32:
+        return hcclFloat;
+    case INFINI_DTYPE_F16:
+        return hcclHalf;
+    default:
+        std::abort();
+        return hcclHalf;
+    }
+}
+
+inline hcclRedOp_t getHcclRedOp(infinicclReduceOp_t op) {
+    switch (op) {
+    case INFINICCL_SUM:
+        return hcclSum;
+    case INFINICCL_PROD:
+        return hcclProd;
+    case INFINICCL_MAX:
+        return hcclMax;
+    case INFINICCL_MIN:
+        return hcclMin;
+    case INFINICCL_AVG:
+        return hcclAvg;
+    default:
+        std::abort();
+        return hcclSum;
+    }
+}
+
+inline hcclComm_t getHcclComm(infinicclComm_t comm) {
+    return static_cast<hcclComm_t>(comm->comm);
+}
+
+namespace infiniccl::maca {
+
+infiniStatus_t commInitAll(
+    infinicclComm_t *comms,
+    int ndevice,
+    const int *device_ids) {
+
+    std::vector<hcclComm_t> hccl_comms(ndevice);
+    CHECK_HCCL(hcclCommInitAll(hccl_comms.data(), ndevice, (int const *)device_ids));
+
+    for (int i = 0; i < ndevice; i++) {
+        comms[i] = new InfinicclComm{INFINI_DEVICE_METAX, device_ids[i], (void *)(hccl_comms[i])};
+    }
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t commDestroy(infinicclComm_t comm) {
+    CHECK_HCCL(hcclCommDestroy(getHcclComm(comm)));
+    delete comm;
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t allReduce(
+    void *sendbuf,
+    void *recvbuf,
+    size_t count,
+    infiniDtype_t datatype,
+    infinicclReduceOp_t op,
+    infinicclComm_t comm,
+    infinirtStream_t stream) {
+
+    if (datatype != INFINI_DTYPE_F32 && datatype != INFINI_DTYPE_F16) {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+
+    CHECK_HCCL(hcclAllReduce(sendbuf, recvbuf, count, getHcclDtype(datatype),
+                             getHcclRedOp(op), getHcclComm(comm), getMacaStream(stream)));
+
+    return INFINI_STATUS_SUCCESS;
+}
+} // namespace infiniccl::maca

src/infiniccl/maca/infiniccl_maca.h

+#ifndef INFINICCL_MACA_H_
+#define INFINICCL_MACA_H_
+
+#include "../infiniccl_impl.h"
+
+#if defined(ENABLE_METAX_API) && defined(ENABLE_CCL)
+INFINICCL_DEVICE_API_IMPL(maca)
+#else
+INFINICCL_DEVICE_API_NOOP(maca)
+#endif
+
+#endif /* INFINICCL_MACA_H_ */

xmake.lua

@@ -245,6 +245,9 @@ target("infiniccl")
     if has_config("ascend-npu") then
         add_deps("infiniccl-ascend")
     end
+    if has_config("metax-gpu") then
+        add_deps("infiniccl-metax")
+    end
     
     set_languages("cxx17")
 

xmake/maca.lua

@@ -47,3 +47,19 @@ target("infinirt-metax")
     add_cxflags("-lstdc++ -fPIC")
     add_files("../src/infinirt/maca/*.cc")
 target_end()
+
+target("infiniccl-metax")
+    set_kind("static")
+    add_deps("infinirt")
+    on_install(function (target) end)
+    set_warnings("all", "error")
+    if not is_plat("windows") then
+        add_cxflags("-fPIC")
+    end
+    if has_config("ccl") then
+        add_links("libhccl.so")
+        add_files("../src/infiniccl/maca/*.cc")
+    end
+    set_languages("cxx17")
+    
+target_end()