[Graph]Add AMX_BF16 and AMX_INT8 kernels (#126)

Signed-off-by: Hengyu Meng <hengyu.meng@intel.com> Co-authored-by: luoyu-intel <yu.luo@intel.com>

Author: airMeng

.github/CODEOWNERS

@@ -26,4 +26,5 @@
 
 
 /intel_extension_for_transformers/backends/neural_engine/Cmake* yu.luo@intel.com
-/intel_extension_for_transformers/backends/neural_engine/cmake/ yu.luo@intel.com
+/intel_extension_for_transformers/backends/neural_engine/cmake/ yu.luo@intel.com
+/intel_extension_for_transformers/backends/neural_engine/graph/jblas yu.luo@intel.com

intel_extension_for_transformers/backends/neural_engine/graph/CMakeLists.txt

@@ -74,6 +74,14 @@ option(NE_PROFILING "neural_engine: use Profiling"
 if (NE_PROFILING)
  add_compile_definitions(NE_PERF)
 endif()
+option(NE_GELU_VEC "neural_engine: enable vec in gelu" ON)
+if (NE_GELU_VEC)
+ add_compile_definitions(NE_GELU_USE_VEC)
+endif()
+
+if(NE_BUILD_TESTS)
+ enable_testing()
+endif()
 
 if (MSVC)
  add_compile_definitions(_CRT_SECURE_NO_WARNINGS NOMINMAX)

intel_extension_for_transformers/backends/neural_engine/graph/CMakePresets.json

@@ -23,21 +23,6 @@
  "inherits": "linux-debug",
  "cacheVariables": { "CMAKE_BUILD_TYPE": "Release" }
  },
- {
- "name": "macos-debug",
- "displayName": "macOS Debug",
- "description": "Target a remote macOS system.",
- "generator": "Ninja",
- "binaryDir": "${sourceDir}/out/build/${presetName}",
- "installDir": "${sourceDir}/out/install/${presetName}",
- "cacheVariables": { "CMAKE_BUILD_TYPE": "Debug" },
- "condition": {
- "type": "equals",
- "lhs": "${hostSystemName}",
- "rhs": "Darwin"
- },
- "vendor": { "microsoft.com/VisualStudioRemoteSettings/CMake/1.0": { "sourceDir": "$env{HOME}/.vs/$ms{projectDirName}" } }
- },
  {
  "name": "windows-base",
  "description": "Target Windows with the Visual Studio development environment.",
@@ -72,24 +57,6 @@
  "description": "Target Windows (64-bit) with the Visual Studio development environment. (RelWithDebInfo)",
  "inherits": "x64-debug",
  "cacheVariables": { "CMAKE_BUILD_TYPE": "Release" }
- },
- {
- "name": "x86-debug",
- "displayName": "x86 Debug",
- "description": "Target Windows (32-bit) with the Visual Studio development environment. (Debug)",
- "inherits": "windows-base",
- "architecture": {
- "value": "x86",
- "strategy": "external"
- },
- "cacheVariables": { "CMAKE_BUILD_TYPE": "Debug" }
- },
- {
- "name": "x86-release",
- "displayName": "x86 Release",
- "description": "Target Windows (32-bit) with the Visual Studio development environment. (RelWithDebInfo)",
- "inherits": "x86-debug",
- "cacheVariables": { "CMAKE_BUILD_TYPE": "Release" }
  }
  ]
 }

intel_extension_for_transformers/backends/neural_engine/graph/application/ChatGPTJ/pybind_gptj.cpp

@@ -33,9 +33,9 @@ class gptj_quant_layer : public quant_layer_base {
  virtual quant_params_internal get_layer_config(std::string layername, std::vector<int64_t> ne,
  ne_type type) override {
  bool quantize = layername.rfind("weight") == layername.size() - 6; // ends with 'weight'?
- if (layername == "transformer.wte.weight") {
+ if (layername == "transformer.wte.weight" || layername == "lm_head.weight") {
  // special layer process, can be loaded by config file
- return quant_params_internal(); // return q4_0 to cover the usage of getrow
+ return quant_params_internal{quant_bits::count}; // skip for head and tail layers
  }
  quantize &= (ne.size() == 2);
  if (quantize) {

intel_extension_for_transformers/backends/neural_engine/graph/application/ChatGPTJ/quant_gptj.cpp

@@ -786,21 +786,13 @@ size_t jblas_quantize(const float* f32ptr, void* dstpr, const quant_params param
  if (params.compute_type == "int8") {
  using GemmKernel = jblas::wrapper::gemm_default::weight_comp::avx512_vnni::GemmKernelDynamicQuantS4KBlock;
  static GemmKernel kernel;
- if (cd->AVX512F()) {
- packedw =
- kernel.getWeightPtr()->compressWeightTranspose<JblasAVX512F>(n, k, f32ptr, k, params.block_size, type);
- } else {
- packedw = kernel.getWeightPtr()->compressWeightTranspose<JblasNoSIMD>(n, k, f32ptr, k, params.block_size, type);
- }
+ assert(cd->AVX512F());
+ packedw = kernel.getWeightPtr()->compressWeightTranspose(n, k, f32ptr, k, params.block_size, type);
  } else if (params.compute_type == "fp32") {
  using GemmKernel = jblas::wrapper::gemm_default::weight_comp::avx512f::GemmKernelS4KBlock;
  static GemmKernel kernel;
- if (cd->AVX512F()) {
- packedw =
- kernel.getWeightPtr()->compressWeightTranspose<JblasAVX512F>(n, k, f32ptr, k, params.block_size, type);
- } else {
- packedw = kernel.getWeightPtr()->compressWeightTranspose<JblasNoSIMD>(n, k, f32ptr, k, params.block_size, type);
- }
+ assert(cd->AVX512F());
+ packedw = kernel.getWeightPtr()->compressWeightTranspose(n, k, f32ptr, k, params.block_size, type);
  }
  } else if (params.bits == 8) {
  // TODO add 8bit quantization

intel_extension_for_transformers/backends/neural_engine/graph/application/common.cpp

@@ -26,3 +26,32 @@ if(NOT WIN32)
  target_link_libraries(ne_layers PUBLIC rt)
 endif()
 
+add_compile_definitions(NE_USE_RN_BF16FP16=1)
+
+
+if (NE_BUILD_TESTS)
+
+function(add_test_target src)
+ get_filename_component(test_target ${src} NAME_WE)
+ get_filename_component(src_dir ${src} DIRECTORY)
+ string(REGEX REPLACE [/\\] "_" src_dir ${src_dir})
+ if(src_dir)
+ set (test_target "${src_dir}_${test_target}")
+ endif()
+ set (test_target "test_${test_target}")
+ add_executable_w_warning(${test_target} ${src})
+ target_compile_definitions(${test_target} PRIVATE NE_TESTS)
+ target_compile_options(${test_target} PRIVATE -fsanitize=address)
+ target_link_options(${test_target} PRIVATE -fsanitize=address)
+ target_include_directories(${test_target} PUBLIC .)
+ target_link_libraries(${test_target} PUBLIC Threads::Threads jblas::jblas ne_vec)
+ if(NOT WIN32)
+ target_link_libraries(${test_target} PUBLIC rt)
+ endif()
+ add_test(NAME ${test_target} COMMAND ${test_target})
+ set_tests_properties(${test_target} PROPERTIES LABELS "${src_dir}_test")
+endfunction()
+
+add_test_target(layers/mha_dense.cpp)
+
+endif()

intel_extension_for_transformers/backends/neural_engine/graph/core/CMakeLists.txt

@@ -13,10 +13,11 @@
 // limitations under the License.
 #pragma once
 
-#include <stdint.h>
-#include <stddef.h>
-#include <stdbool.h>
 #include <assert.h>
+#include <math.h>
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
 #include <string.h>
 
 #ifdef __cplusplus
@@ -26,6 +27,7 @@ extern "C" {
 // floating point type used to accumulate sums
 typedef double ne_float;
 typedef uint16_t ne_fp16_t;
+typedef uint16_t ne_bf16_t;
 
 enum ne_type {
  NE_TYPE_F32 = 0,

intel_extension_for_transformers/backends/neural_engine/graph/core/data_types.h

@@ -47,7 +47,7 @@ enum ne_op {
  NE_OP_RMS_NORM_BACK,
 
  NE_OP_MUL_MAT,
-
+ NE_OP_MUL_MAT_BIAS,
  NE_OP_SCALE,
  NE_OP_SET,
  NE_OP_CPY,
@@ -72,6 +72,8 @@ enum ne_op {
  // LLM related
  NE_OP_MUL_QKV,
  NE_OP_MUL_FFN_SILU,
+ NE_OP_MUL_FFN_GELU,
+ NE_OP_MUL_FFN_ADD_GELU,
  NE_OP_FLASH_ATTN,
  NE_OP_FLASH_FF,
 

intel_extension_for_transformers/backends/neural_engine/graph/core/layers/Ops.h

@@ -40,6 +40,14 @@ inline static void ne_vec_set_f16(const int n, ne_fp16_t* x, const int32_t v) {
  for (int i = 0; i < n; ++i) x[i] = v;
 }
 
+inline static void ne_vec_srl_i32(const int n, int32_t* z, const int32_t* x, int32_t v) {
+ for (int i = 0; i < n; ++i) z[i] = x[i] >> v;
+}
+
+inline static void ne_vec_and_i32(const int n, int32_t* z, const int32_t* x, const int32_t* y) {
+ for (int i = 0; i < n; ++i) z[i] = x[i] & y[i];
+}
+
 inline static void ne_vec_add_f32(const int n, float* z, const float* x, const float* y) {
  for (int i = 0; i < n; ++i) z[i] = x[i] + y[i];
 }
@@ -171,22 +179,38 @@ inline static void ne_vec_gelu_f16(const int n, ne_fp16_t* y, const ne_fp16_t* x
  }
 }
 
-#ifdef NE_GELU_FP16
-inline static void ne_vec_gelu_f32(const int n, float* y, const float* x) {
- uint16_t t;
- for (int i = 0; i < n; ++i) {
- ne_fp16_t fp16 = NE_FP32_TO_FP16(x[i]);
- memcpy(&t, &fp16, sizeof(uint16_t));
- y[i] = NE_FP16_TO_FP32(table_gelu_f16[t]);
- }
+inline static void ne_vec_tanh_f32(const int n, float* y, const float* x) {
+ for (int i = 0; i < n; i++) y[i] = tanhf(x[i]);
 }
-#else
+
 inline static void ne_vec_gelu_f32(const int n, float* y, const float* x) {
+#ifdef NE_GELU_USE_VEC
+ // compute G(x) = sqrt_root_two_over_pi * x * (1 + fitting_const * x * x)
+ float* aux0 = (float*)malloc(n * sizeof(float));
+ ne_vec_sqr_f32(n, aux0, x);
+ float* aux1 = (float*)malloc(n * sizeof(float));
+ ne_vec_set_f32(n, aux1, 1.0f);
+ ne_vec_mad_f32(n, aux1, aux0, GELU_COEF_A);
+ ne_vec_mul_f32(n, aux0, x, aux1);
+ ne_vec_set_f32(n, aux1, SQRT_2_OVER_PI);
+ ne_vec_mul_f32(n, aux1, aux0, aux1);
+
+ // compute tanh(G(x))
+ ne_vec_tanh_f32(n, aux0, aux1);
+ // Gelu(x)= 0.5f * x * (1.0f + tanh(G(x)))
+ ne_vec_acc1_f32(n, aux0, 1.0f);
+ ne_vec_mul_f32(n, y, x, aux0);
+ ne_vec_set_f32(n, aux0, 0.5f);
+ ne_vec_mul_f32(n, y, y, aux0);
+
+ free(aux0);
+ free(aux1);
+#else
  for (int i = 0; i < n; ++i) {
  y[i] = ne_gelu_f32(x[i]);
  }
-}
 #endif
+}
 
 // Sigmoid Linear Unit (SiLU) function
 inline static float ne_silu_f32(float x) { return x / (1.0f + expf(-x)); }