From 51df132d876e933581ff9f537b16757f4e055f74 Mon Sep 17 00:00:00 2001 From: liyinghui-qy Date: Mon, 6 Jan 2025 17:00:10 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0CPU=20AllReduceSum=E7=9A=84?= =?UTF-8?q?=E6=94=AF=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/ccl/cpu/infiniccl_cpu.cc | 43 ++++++++++++++++++++++++++++++++++++ src/ccl/cpu/infiniccl_cpu.h | 15 +++++++++++++ src/ccl/infiniccl.cc | 8 +++++++ xmake.lua | 33 +++++++++++++++++++++++++++ 4 files changed, 99 insertions(+) create mode 100644 src/ccl/cpu/infiniccl_cpu.cc create mode 100644 src/ccl/cpu/infiniccl_cpu.h diff --git a/src/ccl/cpu/infiniccl_cpu.cc b/src/ccl/cpu/infiniccl_cpu.cc new file mode 100644 index 0000000..1d57e25 --- /dev/null +++ b/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; +} \ No newline at end of file diff --git a/src/ccl/cpu/infiniccl_cpu.h b/src/ccl/cpu/infiniccl_cpu.h new file mode 100644 index 0000000..a5f9bcf --- /dev/null +++ b/src/ccl/cpu/infiniccl_cpu.h @@ -0,0 +1,15 @@ + +#ifndef INFINICCL_CPU_H_ +#define INFINICCL_CPU_H_ +#include +#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_ */ \ No newline at end of file diff --git a/src/ccl/infiniccl.cc b/src/ccl/infiniccl.cc index dcc2604..1928ef9 100644 --- a/src/ccl/infiniccl.cc +++ b/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; } diff --git a/xmake.lua b/xmake.lua index b561f88..1d86267 100644 --- a/xmake.lua +++ b/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