增加CPU AllReduceSum的支持

src/ccl/cpu/infiniccl_cpu.cc

@@ -0,0 +1,43 @@
+#include "infiniccl_cpu.h"
+
+inline MPI_Datatype getCpuDtype(InfiniDataType_t datatype) {
+    MPI_Datatype MPI_HALF;
+    MPI_Type_contiguous(2, MPI_BYTE, &MPI_HALF);
+    MPI_Type_commit(&MPI_HALF);
+    switch (datatype) {
+    case INFINI_F32:
+        return MPI_FLOAT;
+    case INFINI_F16:
+        return MPI_HALF;
+    default:
+        return MPI_HALF;
+    }
+}
+
+infinicclStatus_t infinicclCpuCommInitAll(infinicclComm_t *comms) {
+    MPI_Init(NULL, NULL);
+    comms[0] = new InfiniComm{DEVICE_CPU, (unsigned int)0, (void*) MPI_COMM_WORLD};
+    return INFINICCL_STATUS_SUCCESS;
+}
+
+infinicclStatus_t infinicclCpuCommDestroy(infinicclComm_t comm) {
+    delete comm;
+    return INFINICCL_STATUS_SUCCESS;
+}
+
+infinicclStatus_t infinicclCpuAllReduceSum(infinicclComm_t comms, void *sendbuf,
+                              void *recvbuf, size_t count,
+                              InfiniDataType_t datatype) {
+    if (datatype != INFINI_F32 && datatype != INFINI_F16) {
+        return INFINICCL_STATUS_BAD_DATATYPE;
+    }
+
+    MPI_Datatype MPI_HALF;
+    MPI_Type_contiguous(2, MPI_BYTE, &MPI_HALF);
+    MPI_Type_commit(&MPI_HALF);
+    MPI_Datatype datatype_cpu = getCpuDtype(datatype);
+    MPI_Comm* comm = (MPI_Comm*)comms[0].comm;
+    MPI_Allreduce(sendbuf, recvbuf, count, datatype_cpu, MPI_SUM, *comm);
+
+    return INFINICCL_STATUS_SUCCESS;
+}

src/ccl/cpu/infiniccl_cpu.h

@@ -0,0 +1,15 @@
+
+#ifndef INFINICCL_CPU_H_
+#define INFINICCL_CPU_H_
+#include <mpi.h>
+#include "infiniccl.h"
+
+infinicclStatus_t infinicclCpuCommInitAll(infinicclComm_t *comms);
+
+infinicclStatus_t infinicclCpuCommDestroy(infinicclComm_t comm);
+
+infinicclStatus_t infinicclCpuAllReduceSum(infinicclComm_t comm, void *sendbuf,
+                              void *recvbuf, size_t count,
+                              InfiniDataType_t datatype);
+
+#endif /* INFINICCL_CPU_H_ */

src/ccl/infiniccl.cc

@@ -2,6 +2,7 @@
 #include "../runtime/runtime.h"
 #include "./ascend/infiniccl_ascend.h"
 #include "./cuda/infiniccl_cuda.h"
+#include "./cpu/infiniccl_cpu.h"
 
 __C infinicclStatus_t infinicclCommInitAll(DeviceType deviceType,
                                            infinicclComm_t *comms,
@@ -11,6 +12,8 @@ __C infinicclStatus_t infinicclCommInitAll(DeviceType deviceType,
         return infinicclCudaCommInitAll(comms, numDevices, deviceIDs);
     case DEVICE_ASCEND:
         return infinicclAscendCommInitAll(comms, numDevices, deviceIDs);
+    case DEVICE_CPU:
+        return infinicclCpuCommInitAll(comms);
     default:
         return INFINICCL_STATUS_DEVICE_NOT_SUPPORTED;
     }
@@ -25,6 +28,8 @@ __C infinicclStatus_t infinicclCommDestroy(infinicclComm_t comm) {
         return infinicclCudaCommDestroy(comm);
     case DEVICE_ASCEND:
         return infinicclAscendCommDestroy(comm);
+    case DEVICE_CPU:
+        return infinicclCpuCommDestroy(comm);
     default:
         return INFINICCL_STATUS_DEVICE_NOT_SUPPORTED;
     }
@@ -47,6 +52,9 @@ __C infinicclStatus_t infinicclAllReduceSum(infinicclComm_t comm, void *sendbuf,
     case DEVICE_ASCEND:
         return infinicclAscendAllReduceSum(comm, sendbuf, recvbuf, count,
                                            datatype, stream);
+    case DEVICE_CPU:
+        return infinicclCpuAllReduceSum(comm, sendbuf, recvbuf, count,
+                                           datatype);
     default:
         return INFINICCL_STATUS_DEVICE_NOT_SUPPORTED;
     }

xmake.lua

@@ -8,6 +8,13 @@ option("omp")
     set_description("Enable or disable OpenMP support")
 option_end()
 
+option("cpu")
+    set_default(false)
+    set_showmenu(true)
+    set_description("Enable or disable CPU functions")
+    add_defines("ENABLE_CPU")
+option_end()
+
 option("nv-gpu")
     set_default(false)
     set_showmenu(true)
@@ -52,6 +59,26 @@ if has_config("infer") then
     add_includedirs(infini_root .. "/include")
 end
 
+if has_config("cpu") then
+    add_defines("ENABLE_CPU")
+    target("cpu")
+        set_kind("static")
+        if not is_plat("windows") then
+            add_cxflags("-fPIC")
+        end
+        set_languages("cxx17")
+        add_cxflags("-fopenmp")
+        add_ldflags("-fopenmp")
+        if has_config("ccl") then
+            add_includedirs("$(env MPI_HOME)/include")
+            add_linkdirs("$(env MPI_HOME)/lib")
+            add_packages("mpi")
+            add_links("mpi")
+            add_files("src/ccl/cpu/*.cc")
+        end
+    target_end()
+end
+
 if has_config("nv-gpu") then
     add_defines("ENABLE_NV_GPU")
     target("nv-gpu")
@@ -142,6 +169,9 @@ if has_config("ccl") then
 target("infiniccl")
     set_kind("shared")
     add_deps("infinirt")
+    if has_config("cpu") then
+        add_deps("cpu")
+    end
     if has_config("nv-gpu") then
         add_deps("nv-gpu")
     end
@@ -178,6 +208,9 @@ target("infini_infer_test")
     set_languages("cxx17")
     on_install(function (target) end)
     add_includedirs("src")
+    if has_config("cpu") then
+        add_deps("cpu")
+    end
     if has_config("nv-gpu") then
         add_deps("nv-gpu")
     end