Merge branch 'dev' into dev-hardware

Author: kilinchange

src/02hardware/CMakeLists.txt

@@ -2,17 +2,15 @@ cmake_minimum_required(VERSION 3.12 FATAL_ERROR)
 project(hardware VERSION 0.0.0 LANGUAGES CXX)
 message(STATUS "Project " ${PROJECT_NAME} " version " ${PROJECT_VERSION})
 
-# Source files
 file(GLOB_RECURSE HARDWARE_SRC src/*.cc src/*.cpp)
+add_library(hardware STATIC ${HARDWARE_SRC})
+target_link_libraries(hardware PUBLIC common)
+target_include_directories(hardware PUBLIC include)
 
 if(USE_CUDA)
-    file(GLOB_RECURSE HARDWARE_CUDA_SRC src/devices/nvidia/*.cu)
+    target_include_directories(hardware PUBLIC ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
 endif()
 
-add_library(hardware STATIC ${HARDWARE_SRC} ${HARDWARE_CUDA_SRC} ${HARDWARE_BANG_SRC})
-target_link_libraries(hardware PUBLIC common)
-target_include_directories(hardware PUBLIC include)
-
 file(GLOB_RECURSE HARDWARE_TEST test/*.cpp)
 if(HARDWARE_TEST)
     add_executable(hardware_test ${HARDWARE_TEST})

src/02hardware/include/hardware/device.h

@@ -52,7 +52,7 @@ namespace refactor::hardware {
 
         virtual ~Device() = default;
         virtual Type type() const noexcept = 0;
-        virtual void setContext() const noexcept;
+        virtual void setContext() const;
 
         Arc<Blob> malloc(size_t);
         Arc<Blob> absorb(Arc<Blob> &&);

src/02hardware/include/hardware/devices/nvidia.h

@@ -8,7 +8,7 @@ namespace refactor::hardware {
     class Nvidia final : public Device {
     public:
         explicit Nvidia(int32_t card);
-        void setContext() const noexcept final;
+        void setContext() const final;
         Type type() const noexcept final {
             return Type::Nvidia;
         }

src/02hardware/src/device.cc

@@ -56,7 +56,7 @@ namespace refactor::hardware {
     Device::Device(decltype(_card) card, decltype(_mem) mem)
         : _card(card), _mem(std::move(mem)) {}
 
-    void Device::setContext() const noexcept {}
+    void Device::setContext() const {}
     auto Device::malloc(size_t size) -> Arc<Blob> {
         return Arc<Blob>(new Blob(this, size));
     }

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

@@ -1,15 +1,29 @@
 #include "functions.hh"
 #include "hardware/devices/mlu.h"
 #include "hardware/mem_pool.h"
+
+#ifdef USE_BANG
+#include "cnrt.h"
 #include "memory.hh"
 
+#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 {
 
     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();
+        unsigned deviceCount;
+        BANG_ASSERT(cnrtGetDeviceCount(&deviceCount));
+        ASSERT(0 <= card && card < deviceCount, "Invalid card id: {}", card);
+        BANG_ASSERT(cnrtSetDevice(card));
+
+        size_t free, total;
+        BANG_ASSERT(cnrtMemGetInfo(&free, &total));
         auto size = std::min(free, std::max(5ul << 30, total * 4 / 5));
         fmt::println("initializing Cambricon MLU {}, memory {} / {}, alloc {}",
                      card, free, total, size);
@@ -25,7 +39,9 @@ namespace refactor::hardware {
     Mlu::Mlu(int32_t card) : Device(card, bangMemory(card)) {}
 
     void Mlu::setContext() const noexcept {
-        setDevice(_card);
+#ifdef USE_BANG
+        BANG_ASSERT(cnrtSetDevice(_card));
+#endif
     }
 
 }// namespace refactor::hardware

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

@@ -1,8 +1,17 @@
+#ifdef USE_BANG
+
 #include "memory.hh"
-#include "functions.hh"
+#include "cnrt.h"
+#include "common.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));           \
+    }
 
 namespace refactor::hardware {
-#ifdef USE_BANG
+
     using M = MluMemory;
 
     void *M::malloc(size_t size) {
@@ -28,6 +37,6 @@ namespace refactor::hardware {
                                CNRT_MEM_TRANS_DIR_PEER2PEER));
         return dst;
     }
-#endif
 
 }// namespace refactor::hardware
+#endif

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

@@ -1,31 +1,50 @@
 #include "functions.cuh"
 #include "hardware/devices/nvidia.h"
 #include "hardware/mem_pool.h"
-#include "memory.cuh"
+
+#ifdef USE_CUDA
+#include "memory.hh"
+#include <cuda_runtime.h>
+
+#define CUDA_ASSERT(STATUS)                                                          \
+    if (auto status = (STATUS); status != cudaSuccess) {                             \
+        RUNTIME_ERROR(fmt::format("cuda failed on \"" #STATUS "\" with \"{}\" ({})", \
+                                  cudaGetErrorString(status), (int) status));        \
+    }
+#endif
 
 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();
+        int deviceCount;
+        CUDA_ASSERT(cudaGetDeviceCount(&deviceCount));
+        ASSERT(0 <= card && card < deviceCount, "Invalid card id: {}", card);
+        CUDA_ASSERT(cudaSetDevice(card));
+
+        size_t free, total;
+        CUDA_ASSERT(cudaMemGetInfo(&free, &total));
         auto size = std::min(free, std::max(5ul << 30, total * 4 / 5));
-        fmt::println("initializing Nvidia GPU {}, memory {} / {}, alloc {}",
-                     card, free, total, size);
+        cudaDeviceProp prop;
+        CUDA_ASSERT(cudaGetDeviceProperties(&prop, 0));
+        size_t alignment = prop.textureAlignment;
+        fmt::println("initializing Nvidia GPU {}, memory {} / {}, alloc {}, alignment {}",
+                     card, free, total, size, alignment);
         return std::make_shared<MemPool>(
             std::make_shared<NvidiaMemory>(),
             size,
-            256ul);
+            alignment);
 #else
         return nullptr;
 #endif
     }
 
     Nvidia::Nvidia(int32_t card) : Device(card, cudaMemory(card)) {}
 
-    void Nvidia::setContext() const noexcept {
-        setDevice(_card);
+    void Nvidia::setContext() const {
+#ifdef USE_CUDA
+        CUDA_ASSERT(cudaSetDevice(_card));
+#endif
     }
 
 }// namespace refactor::hardware