feat: 接入CNNL，并添加unary/binary/softmax/batchnorm/reduce/transpose/pooling算子

Author: Chamberlain0w0

src/02hardware/CMakeLists.txt

@@ -3,14 +3,10 @@ project(hardware VERSION 0.0.0 LANGUAGES CXX)
 message(STATUS "Project " ${PROJECT_NAME} " version " ${PROJECT_VERSION})
 
 # Source files
-file(GLOB HARDWARE_SRC src/*.cc src/*.cpp src/devices/cpu/*.cc)
+file(GLOB_RECURSE HARDWARE_SRC src/*.cc src/*.cpp)
 
 if(USE_CUDA)
-    file(GLOB_RECURSE HARDWARE_CUDA_SRC src/devices/nvidia/*.cu src/devices/nvidia/*.cc)
-endif()
-
-if(USE_BANG)
-    file(GLOB_RECURSE HARDWARE_BANG_SRC src/devices/mlu/*.cc)
+    file(GLOB_RECURSE HARDWARE_CUDA_SRC src/devices/nvidia/*.cu)
 endif()
 
 add_library(hardware STATIC ${HARDWARE_SRC} ${HARDWARE_CUDA_SRC} ${HARDWARE_BANG_SRC})

src/02hardware/src/device_manager.cpp

@@ -1,6 +1,7 @@
 #include "hardware/device_manager.h"
 #include "hardware/devices/cpu.h"
 #include "hardware/devices/nvidia.h"
+#include "hardware/devices/mlu.h"
 
 namespace refactor::hardware::device {
 
@@ -37,6 +38,7 @@ namespace refactor::hardware::device {
         using T = Device::Type;
         // clang-format off
         auto device = type == T::Nvidia ? std::make_shared<Nvidia>(card)
+                    : type == T::Mlu    ? std::make_shared<Mlu>(card)
                     : UNREACHABLEX(Arc<Device>, "");
         // clang-format on
         auto [kind, ok] = DEVICES.try_emplace(static_cast<int32_t>(type));

src/02hardware/src/devices/mlu/device.cc

@@ -1,21 +1,25 @@
-#include "functions.cc"
+#include "functions.hh"
 #include "hardware/devices/mlu.h"
 #include "hardware/mem_pool.h"
 #include "memory.hh"
 
 namespace refactor::hardware {
 
     static Arc<Memory> bangMemory(int32_t card) {
+#ifdef USE_BANG
         ASSERT(0 <= card && card < getDeviceCount(), "Invalid card id: {}", card);
         setDevice(card);
         auto [free, total] = getMemInfo();
         auto size = std::min(free, std::max(5ul << 30, total * 4 / 5));
-        fmt::println("initializing Nvidia GPU {}, memory {} / {}, alloc {}",
+        fmt::println("initializing Cambricon MLU {}, memory {} / {}, alloc {}",
                      card, free, total, size);
         return std::make_shared<MemPool>(
             std::make_shared<MluMemory>(),
             size,
             256ul);
+#else
+        return nullptr;
+#endif
     }
 
     Mlu::Mlu(int32_t card) : Device(card, bangMemory(card)) {}

src/02hardware/src/devices/mlu/functions.cc

@@ -2,6 +2,7 @@
 
 namespace refactor::hardware {
 
+#ifdef USE_BANG
     int getDeviceCount() {
         unsigned deviceCount;
         BANG_ASSERT(cnrtGetDeviceCount(&deviceCount));
@@ -15,5 +16,6 @@ namespace refactor::hardware {
         BANG_ASSERT(cnrtMemGetInfo(&memInfo.free, &memInfo.total));
         return memInfo;
     }
+#endif
 
 }// namespace refactor::hardware

src/02hardware/src/devices/mlu/functions.hh

@@ -1,14 +1,17 @@
 #ifndef HARDWARE_DEVICES_MLU_FUNCTIONS_CUH
 #define HARDWARE_DEVICES_MLU_FUNCTIONS_CUH
 
-#include "cnrt.h"
 #include "common.h"
 
+#ifdef USE_BANG
+#include "cnrt.h"
+
 #define BANG_ASSERT(STATUS)                                                          \
     if (auto status = (STATUS); status != CNRT_RET_SUCCESS) {                        \
         RUNTIME_ERROR(fmt::format("bang failed on \"" #STATUS "\" with \"{}\" ({})", \
                                   cnrtGetErrorStr(status), (int) status));           \
     }
+#endif 
 
 namespace refactor::hardware {
 

src/02hardware/src/devices/mlu/memory.cc

@@ -2,6 +2,7 @@
 #include "functions.hh"
 
 namespace refactor::hardware {
+#ifdef USE_BANG
     using M = MluMemory;
 
     void *M::malloc(size_t size) {
@@ -27,5 +28,6 @@ namespace refactor::hardware {
                                CNRT_MEM_TRANS_DIR_PEER2PEER));
         return dst;
     }
+#endif
 
 }// namespace refactor::hardware

src/02hardware/src/devices/nvidia/device.cc

@@ -6,6 +6,7 @@
 namespace refactor::hardware {
 
     static Arc<Memory> cudaMemory(int32_t card) {
+#ifdef USE_CUDA
         ASSERT(0 <= card && card < getDeviceCount(), "Invalid card id: {}", card);
         setDevice(card);
         auto [free, total] = getMemInfo();
@@ -16,6 +17,9 @@ namespace refactor::hardware {
             std::make_shared<NvidiaMemory>(),
             size,
             256ul);
+#else
+        return nullptr;
+#endif
     }
 
     Nvidia::Nvidia(int32_t card) : Device(card, cudaMemory(card)) {}

src/04kernel/src/collectors/batch_normalization.cc

@@ -1,6 +1,7 @@
 #include "kernel/collectors/batch_normalization.h"
 #include "../kernels/batch_normalization/cpu_kernel.hh"
 #include "../kernels/batch_normalization/cudnn_kernel.hh"
+#include "../kernels/batch_normalization/cnnl_kernel.hh"
 
 namespace refactor::kernel {
 
@@ -20,6 +21,9 @@ namespace refactor::kernel {
             case decltype(_target)::Nvidia:
                 REGISTER(BatchNormalizationCudnn)
                 break;
+            case decltype(_target)::Mlu:
+                REGISTER(BatchNormalizationCnnl)
+                break;
             default:
                 UNREACHABLEX(void, "Unknown target");
         }

src/04kernel/src/collectors/reduce.cc

@@ -1,6 +1,7 @@
 #include "kernel/collectors/reduce.h"
 #include "../kernels/reduce/cpu_kernel.hh"
 #include "../kernels/reduce/cudnn_kernel.hh"
+#include "../kernels/reduce/cnnl_kernel.hh"
 
 namespace refactor::kernel {
 
@@ -27,6 +28,9 @@ namespace refactor::kernel {
             case decltype(_target)::Nvidia:
                 REGISTER(ReduceCudnn)
                 break;
+            case decltype(_target)::Mlu:
+                REGISTER(ReduceCnnl)
+                break;
             default:
                 UNREACHABLEX(void, "Unknown target");
         }

src/04kernel/src/collectors/simple_binary.cc

@@ -2,6 +2,7 @@
 #include "../kernels/simple_binary/binary_cudnn.hh"
 #include "../kernels/simple_binary/cpu_kernel.hh"
 #include "../kernels/simple_binary/cuda_kernel.hh"
+#include "../kernels/simple_binary/binary_cnnl.hh"
 
 namespace refactor::kernel {
 
@@ -48,6 +49,9 @@ namespace refactor::kernel {
                 REGISTER_BROCAST(BinaryCudnn)
                 REGISTER(BinaryCuda)
                 break;
+            case decltype(_target)::Mlu:
+                REGISTER_BROCAST(BinaryCnnl)
+                break;
             default:
                 UNREACHABLEX(void, "Unknown target");
         }