feat: add Cambricon RMSNorm (#19)

* feat: Add RMSNorm op in cambricon backend.

* refactor: make `Cast` utility to use `Device::Type` template parameter

* refactor: add `Caster` mixin

* refactor: rename `cast**` to `caster**`

* fix: fix the mlu naming to google c++ naming style

* chore: format files with `clang-format`

* refactor: update CUDA kernels to use `Caster`

* fix: fix rmsnorm dispatch to use one dispatch

---------

Co-authored-by: Jiacheng Huang <huangjiacheng0709@outlook.com>

Author: bitzyz

CMakeLists.txt

@@ -179,7 +179,7 @@ if(WITH_CAMBRICON)
 endif()
 
 # If all other platforms are not enabled, CPU is enabled by default.
-if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_METAX AND NOT WITH_MOORE)
+if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_METAX AND NOT WITH_MOORE AND NOT WITH_CAMBRICON)
     add_compile_definitions(WITH_CPU=1)
 endif()
 

src/CMakeLists.txt

@@ -127,11 +127,48 @@ if(WITH_MOORE)
 endif()
 
 if(WITH_CAMBRICON)
-    target_compile_definitions(infiniops PUBLIC WITH_CAMBRICON=1)
+    file(GLOB_RECURSE CAMBRICON_MLU_SOURCES CONFIGURE_DEPENDS "cambricon/*/*.mlu")
+    find_program(CNCC_COMPILER cncc HINTS "${NEUWARE_HOME}/bin" "$ENV{NEUWARE_HOME}/bin" /usr/local/neuware/bin)
+    if(CNCC_COMPILER)
+        message(STATUS "Found cncc: ${CNCC_COMPILER}")
+        set(MLU_COMPILE_OPTS
+            -c --bang-mlu-arch=mtp_592 -O3 -fPIC -Wall -Werror -std=c++17 -pthread
+            -I${CMAKE_CURRENT_SOURCE_DIR} -I${NEUWARE_HOME}/include
+            -idirafter /usr/local/neuware/lib/clang/11.1.0/include
+        )
+        function(compile_mlu_file src_file)
+            get_filename_component(name ${src_file} NAME_WE)
+            get_filename_component(path ${src_file} DIRECTORY)
+            set(out_file "${CMAKE_CURRENT_BINARY_DIR}/${path}/${name}.o")
+            file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/${path}")
+            add_custom_command(OUTPUT ${out_file}
+                COMMAND ${CNCC_COMPILER} ${MLU_COMPILE_OPTS} -c ${src_file} -o ${out_file}
+                DEPENDS ${src_file}
+                COMMENT "Building MLU kernel: ${src_file}"
+            )
+            set_property(DIRECTORY APPEND PROPERTY CAMBRICON_OBJECTS ${out_file})
+        endfunction()
+        foreach(src ${CAMBRICON_MLU_SOURCES})
+            compile_mlu_file(${src})
+        endforeach()
+        get_directory_property(CAMBRICON_OBJECT_FILES CAMBRICON_OBJECTS)
+        if(CAMBRICON_OBJECT_FILES)
+            target_sources(infiniops PRIVATE ${CAMBRICON_OBJECT_FILES})
+        endif()
+    else()
+        message(WARNING "cncc compiler not found. MLU kernels will not be compiled.")
+    endif()
+    target_compile_definitions(infiniops PRIVATE WITH_CAMBRICON=1)
 
     target_include_directories(infiniops PUBLIC "${NEUWARE_HOME}/include")
     target_link_libraries(infiniops PUBLIC ${CAMBRICON_RUNTIME_LIB} ${CAMBRICON_CNNL_LIB} ${CAMBRICON_CNNL_EXTRA_LIB} ${CAMBRICON_PAPI_LIB})
 
+    if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang")
+        target_compile_options(infiniops PUBLIC 
+            "$<$<COMPILE_LANGUAGE:CXX>:SHELL:-idirafter /usr/local/neuware/lib/clang/11.1.0/include>"
+        )
+    endif()
+
     list(APPEND DEVICE_LIST "cambricon")
 endif()
 

src/base/rms_norm.h

@@ -12,15 +12,17 @@ namespace infini::ops {
 class RmsNorm : public Operator<RmsNorm> {
  public:
   RmsNorm(const Tensor input, const Tensor weight, float eps, Tensor out)
-      : eps_{eps},
+      : input_shape_{input.shape()},
         out_shape_{out.shape()},
-        input_shape_{input.shape()},
-        out_strides_{out.strides()},
         input_strides_{input.strides()},
+        out_strides_{out.strides()},
+        eps_{eps},
         dim_{out.size(-1)},
         ndim_{out.ndim()},
         batch_size_{ndim_ == 2 ? out.size(-2) : out.size(-3)},
-        nhead_{ndim_ == 2 ? 1 : out.size(-2)} {}
+        nhead_{ndim_ == 2 ? 1 : out.size(-2)} {
+    assert(input.dtype() == out.dtype());
+  }
 
   RmsNorm(const Tensor input, const Tensor weight, Tensor out)
       : RmsNorm{input, weight, 1e-6f, out} {}

src/cambricon/common.h

@@ -2,24 +2,74 @@
 #define INFINI_OPS_CAMBRICON_COMMON_H_
 
 #include <cnnl.h>
+#include <cnrt.h>
 
 #include "data_type.h"
+#include "device.h"
+
+#define NRAM_MAX_SIZE (1024 * 240)
+
+#ifdef __BANG__
+
+namespace infini::ops::reduce {
+
+constexpr int batch_size = 128 / sizeof(float);
+
+__mlu_func__ void SumInternal(float* dst, float* src, int max_batch) {
+  const int width = max_batch / batch_size;
+
+  if (width >= 4) {
+    __bang_sumpool(dst, src, batch_size, 1, width, 1, width, 1, 1);
+    __bang_reduce_sum(dst, dst, batch_size);
+  } else {
+    float sum = 0.0f;
+    for (int i = 0; i < max_batch; ++i) {
+      sum += src[i];
+    }
+    dst[0] = sum;
+  }
+}
+
+}  // namespace infini::ops::reduce
+
+#endif  // __BANG__
 
 namespace infini::ops::cnnl_utils {
 
 inline cnnlDataType_t GetDataType(DataType dtype) {
   switch (dtype) {
+    case DataType::kInt8:
+      return CNNL_DTYPE_INT8;
+    case DataType::kUInt8:
+      return CNNL_DTYPE_UINT8;
     case DataType::kInt32:
       return CNNL_DTYPE_INT32;
+    case DataType::kInt64:
+      return CNNL_DTYPE_INT64;
     case DataType::kFloat16:
       return CNNL_DTYPE_HALF;
     case DataType::kFloat32:
       return CNNL_DTYPE_FLOAT;
+    case DataType::kBFloat16:
+      return CNNL_DTYPE_BFLOAT16;
+    case DataType::kFloat64:
+      return CNNL_DTYPE_DOUBLE;
     default:
       return CNNL_DTYPE_INVALID;
   }
 }
 
 }  // namespace infini::ops::cnnl_utils
 
+namespace infini::ops::cnrt_utils {
+
+inline void GetLaunchConfig(const Device& device, int* core_per_cluster,
+                            int* cluster_count) {
+  int device_id = device.index();
+  cnrtDeviceGetAttribute(cluster_count, cnrtAttrClusterCount, device_id);
+  cnrtDeviceGetAttribute(core_per_cluster, cnrtAttrMcorePerCluster, device_id);
+}
+
+}  // namespace infini::ops::cnrt_utils
+
 #endif