issue/300 - InfiniCCL Cambricon

src/infiniccl/cambricon/infiniccl_cambricon.cc

+#include "infiniccl_cambricon.h"
+
+#include "../../utils.h"
+#include <cncl.h>
+#include <cnrt.h>
+#include <iostream>
+#include <vector>
+
+#define CHECK_CNCL(API__) CHECK_INTERNAL(API__, CNCL_RET_SUCCESS)
+
+inline cnrtQueue_t getCambriconStream(infinirtStream_t stream) {
+    if (stream == nullptr) {
+        return (cnrtQueue_t)(0);
+    }
+    return static_cast<cnrtQueue_t>(stream);
+}
+
+inline cnclComm_t getCnclComm(infinicclComm_t comm) {
+    return static_cast<cnclComm_t>(comm->comm);
+}
+
+inline cnclDataType_t getCnclDtype(infiniDtype_t datatype) {
+    switch (datatype) {
+    case INFINI_DTYPE_F32:
+        return cnclFloat32;
+    case INFINI_DTYPE_F16:
+        return cnclFloat16;
+    default:
+        std::cerr << "Unsupported data type: " << datatype << std::endl;
+        std::abort();
+        return cnclFloat16;
+    }
+}
+
+inline cnclReduceOp_t getCnclRedOp(infinicclReduceOp_t op) {
+    switch (op) {
+    case INFINICCL_SUM:
+        return cnclSum;
+    case INFINICCL_PROD:
+        return cnclProd;
+    case INFINICCL_MAX:
+        return cnclMax;
+    case INFINICCL_MIN:
+        return cnclMin;
+    default:
+        std::abort();
+        return cnclSum;
+    }
+}
+
+namespace infiniccl::cambricon {
+
+infiniStatus_t commInitAll(
+    infinicclComm_t *comms,
+    int ndevice,
+    const int *device_ids) {
+
+    std::vector<cnclComm_t> cncl_comms(ndevice);
+    std::vector<int> rank_list(ndevice);
+
+    for (int i = 0; i < ndevice; i++) {
+        rank_list[i] = i;
+        CHECK_INTERNAL(cnrtSetDevice(device_ids[i]), CNRT_RET_SUCCESS);
+    }
+
+    CHECK_CNCL(cnclInitComms(cncl_comms.data(), ndevice,
+                             (int const *)device_ids, rank_list.data(),
+                             ndevice, nullptr));
+
+    for (int i = 0; i < ndevice; i++) {
+        comms[i] = new InfinicclComm{INFINI_DEVICE_CAMBRICON, device_ids[i], (void *)(cncl_comms[i])};
+    }
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t commDestroy(infinicclComm_t comm) {
+    CHECK_CNCL(cnclFreeComm(getCnclComm(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_CNCL(cnclAllReduce(sendbuf, recvbuf, count, getCnclDtype(datatype),
+                             getCnclRedOp(op), getCnclComm(comm),
+                             getCambriconStream(stream)));
+
+    return INFINI_STATUS_SUCCESS;
+}
+} // namespace infiniccl::cambricon

src/infiniccl/cambricon/infiniccl_cambricon.h

+#ifndef INFINICCL_CAMBRICON_H_
+#define INFINICCL_CAMBRICON_H_
+
+#include "../infiniccl_impl.h"
+
+#if defined(ENABLE_CAMBRICON_API) && defined(ENABLE_CCL)
+INFINICCL_DEVICE_API_IMPL(cambricon)
+#else
+INFINICCL_DEVICE_API_NOOP(cambricon)
+#endif
+
+#endif /* INFINICCL_CAMBRICON_H_ */

src/infiniccl/infiniccl.cc

 #include "infiniccl.h"
 
 #include "./ascend/infiniccl_ascend.h"
+#include "./cambricon/infiniccl_cambricon.h"
 #include "./cuda/infiniccl_cuda.h"
 #include "./maca/infiniccl_maca.h"
 
@@ -17,6 +18,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_CAMBRICON, cambricon)
         COMM_INIT_ALL(INFINI_DEVICE_METAX, maca)
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -37,6 +39,7 @@ __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_CAMBRICON, cambricon)
         COMM_DESTROY(INFINI_DEVICE_METAX, maca)
 
     default:
@@ -65,6 +68,7 @@ __C infiniStatus_t infinicclAllReduce(
     switch (comm->device_type) {
         ALL_REDUCE(INFINI_DEVICE_NVIDIA, cuda)
         ALL_REDUCE(INFINI_DEVICE_ASCEND, ascend)
+        ALL_REDUCE(INFINI_DEVICE_CAMBRICON, cambricon)
         ALL_REDUCE(INFINI_DEVICE_METAX, maca)
 
     default:

src/infinirt/bang/infinirt_bang.cc

@@ -8,6 +8,7 @@ namespace infinirt::bang {
 infiniStatus_t getDeviceCount(int *count) {
     unsigned int device_count = static_cast<unsigned int>(*count);
     CHECK_BANGRT(cnrtGetDeviceCount(&device_count));
+    *count = static_cast<int>(device_count);
     return INFINI_STATUS_SUCCESS;
 }
 

xmake.lua

@@ -264,6 +264,9 @@ target("infiniccl")
     if has_config("ascend-npu") then
         add_deps("infiniccl-ascend")
     end
+    if has_config("cambricon-mlu") then
+        add_deps("infiniccl-cambricon")
+    end
     if has_config("metax-gpu") then
         add_deps("infiniccl-metax")
     end

xmake/bang.lua

@@ -60,3 +60,38 @@ target("infinirt-cambricon")
     add_files("../src/infinirt/bang/*.cc")
     add_cxflags("-lstdc++ -Wall -Werror -fPIC")
 target_end()
+
+target("infiniccl-cambricon")
+    set_kind("static")
+    add_deps("infinirt")
+    add_deps("infini-utils")
+    set_warnings("all", "error")
+    set_languages("cxx17")
+    on_install(function (target) end)
+    
+    if has_config("ccl") then
+        if is_plat("linux") then
+            add_includedirs(NEUWARE_HOME .. "/include")
+            add_linkdirs(NEUWARE_HOME .. "/lib64")
+            add_links("cncl", "cnrt")
+
+            if has_package("libibverbs") then
+                add_links("ibverbs")
+                add_defines("CNCL_RDMA_ENABLED=1")
+            end
+
+            if is_arch("arm64") then
+                add_defines("CNCL_ARM64_COMPAT_MODE=1")
+            end
+
+            add_rpathdirs(NEUWARE_HOME .. "/lib64")
+            add_runenvs("LD_LIBRARY_PATH", NEUWARE_HOME .. "/lib64")
+
+            add_files("../src/infiniccl/cambricon/*.cc")
+            add_cxflags("-fPIC")
+            add_ldflags("-fPIC")
+        else
+            print("[Warning] CNCL is currently only supported on Linux")
+        end
+    end
+target_end()