feat: add rmsnorm op on cambricon impl

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/caster_.h

@@ -0,0 +1,98 @@
+#ifndef INFINI_OPS_COMMON_CAMBRICON_CASTER_H_
+#define INFINI_OPS_COMMON_CAMBRICON_CASTER_H_
+
+#include "bang_bf16.h"
+#include "bang_fp16.h"
+#include "caster.h"
+
+namespace infini::ops {
+
+template <>
+struct Caster<Device::Type::kCambricon> {
+  template <typename Dst, typename Src>
+  static Dst Cast(Src&& x) {
+    static_assert(!std::is_reference_v<Dst>,
+                  "`Cast` cannot return reference types");
+
+    using PureSrc = std::remove_cv_t<std::remove_reference_t<Src>>;
+    using PureDst = std::remove_cv_t<std::remove_reference_t<Dst>>;
+
+    if constexpr (std::is_same_v<PureSrc, PureDst>) {
+      return std::forward<Src>(x);
+    } else {
+      return HardwareCast<PureDst>(std::forward<Src>(x), PriorityHigh{});
+    }
+  }
+
+ private:
+  template <typename T>
+  using PureType = std::remove_cv_t<std::remove_reference_t<T>>;
+
+  template <typename T>
+  __host__ __device__ static constexpr float ToFloatHelper(T&& x) {
+    using PureSrc = PureType<T>;
+    if constexpr (IsBFloat16<PureSrc>) {
+      return __bfloat162float__(x);
+    } else if constexpr (IsFP16<PureSrc>) {
+      return __half2float(x);
+    } else {
+      return static_cast<float>(std::forward<T>(x));
+    }
+  }
+
+  template <typename Dst>
+  __host__ __device__ static constexpr Dst FromFloatHelper(float f) {
+    using PureDst = PureType<Dst>;
+    if constexpr (IsBFloat16<PureDst>) {
+      return __float2bfloat16__(f);
+    } else if constexpr (IsFP16<PureDst>) {
+      return __float2half__(f);
+    } else {
+      return static_cast<Dst>(f);
+    }
+  }
+
+  // Priority tags for overload resolution.
+  struct PriorityLow {};
+
+  struct PriorityHigh : PriorityLow {};
+
+  // Fallback: lowest priority. This always matches if nothing else does.
+  template <typename Dst, typename Src>
+  __host__ __device__ static constexpr Dst HardwareCast(Src&& x, PriorityLow) {
+    return FromFloatHelper<Dst>(ToFloatHelper(std::forward<Src>(x)));
+  }
+
+// Usage: `DEFINE_DIRECT_CAST(INTRINSIC, CONDITION)`.
+#define DEFINE_DIRECT_CAST(INTRINSIC, ...)                            \
+  template <typename Dst, typename Src>                               \
+  __host__ __device__ static auto HardwareCast(Src x, PriorityHigh)   \
+      -> std::enable_if_t<(__VA_ARGS__),                              \
+                          decltype(INTRINSIC(std::declval<Src>()))> { \
+    return INTRINSIC(x);                                              \
+  }
+
+  DEFINE_DIRECT_CAST(
+      __bfloat162int_rz__,
+      std::is_same_v<PureType<Dst>, int>&& IsBFloat16<PureType<Src>>)
+  DEFINE_DIRECT_CAST(
+      __bfloat162short_rz__,
+      std::is_same_v<PureType<Dst>, short>&& IsBFloat16<PureType<Src>>)
+  DEFINE_DIRECT_CAST(
+      __int2bfloat16_rn__,
+      IsBFloat16<PureType<Dst>>&& std::is_same_v<PureType<Src>, int>)
+  DEFINE_DIRECT_CAST(__int2half_rn__,
+                     IsFP16<PureType<Dst>>&& std::is_same_v<PureType<Src>, int>)
+  DEFINE_DIRECT_CAST(
+      __float2bfloat16__,
+      IsBFloat16<PureType<Dst>>&& std::is_same_v<PureType<Src>, double>)
+  DEFINE_DIRECT_CAST(
+      __float2half__,
+      IsFP16<PureType<Dst>>&& std::is_same_v<PureType<Src>, double>)
+  DEFINE_DIRECT_CAST(__half, IsFP16<PureType<Dst>>&& IsBFloat16<PureType<Src>>)
+#undef DEFINE_DIRECT_CAST
+};
+
+}  // namespace infini::ops
+
+#endif

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