[SparseLib] softmax int8 using lut (#215)

* benchmark done.

* update ci.

Author: zhewang1-intc

nlp_toolkit/backends/neural_engine/SparseLib/include/interface.hpp

@@ -108,13 +108,6 @@ class transpose_matmul_desc : public kernel_desc_proxy {
   virtual ~transpose_matmul_desc() {}
 };
 
-class postop_desc : public kernel_desc_proxy {
- public:
-  postop_desc() {}
-  explicit postop_desc(const operator_desc& op_desc) : kernel_desc_proxy(op_desc) {}
-  virtual ~postop_desc() {}
-};
-
 class eltwiseop_desc : public kernel_desc_proxy {
  public:
   eltwiseop_desc() {}
@@ -129,6 +122,13 @@ class layernorm_ba_desc : public kernel_desc_proxy {
   virtual ~layernorm_ba_desc() {}
 };
 
+class softmax_desc : public kernel_desc_proxy {
+ public:
+  softmax_desc() {}
+  explicit softmax_desc(const operator_desc& op_desc) : kernel_desc_proxy(op_desc) {}
+  virtual ~softmax_desc() {}
+};
+
 /**
  * @brief Derived proxy class, interfacing to the real/cached sparse_matmul_t.
  */
@@ -146,13 +146,6 @@ class transpose_matmul : public kernel_proxy {
   virtual ~transpose_matmul() {}
 };
 
-class postop : public kernel_proxy {
- public:
-  postop() {}
-  explicit postop(const kernel_desc_proxy& kdp) : kernel_proxy(kdp) {}
-  virtual ~postop() {}
-};
-
 class eltwiseop : public kernel_proxy {
  public:
   eltwiseop() {}
@@ -167,5 +160,12 @@ class layernorm_ba : public kernel_proxy {
   virtual ~layernorm_ba() {}
 };
 
+class softmax : public kernel_proxy {
+ public:
+  softmax() {}
+  explicit softmax(const kernel_desc_proxy& kdp) : kernel_proxy(kdp) {}
+  virtual ~softmax() {}
+};
+
 }  // namespace jd
 #endif  // ENGINE_SPARSELIB_INCLUDE_INTERFACE_HPP_

nlp_toolkit/backends/neural_engine/SparseLib/include/jit_domain/jit_eltwise_injector.hpp

@@ -15,6 +15,7 @@
 #ifndef ENGINE_SPARSELIB_INCLUDE_JIT_DOMAIN_JIT_ELTWISE_INJECTOR_HPP_
 #define ENGINE_SPARSELIB_INCLUDE_JIT_DOMAIN_JIT_ELTWISE_INJECTOR_HPP_
 
+#include <glog/logging.h>
 #include <string>
 #include <vector>
 #include <unordered_map>
@@ -54,10 +55,12 @@ class jit_eltwise_injector {
   void quantize_compute_vector_fwd(const Xbyak::Zmm& zmm_src);
   void dequantize_compute_vector_fwd(const Xbyak::Zmm& zmm_src);
   void linear_compute_vector_fwd(const Xbyak::Zmm& zmm_src);
-  void int8_lut_compute_vector_fwd(const Xbyak::Zmm& zmm_src);
+  void bit8_lut_compute_vector_fwd(const Xbyak::Zmm& zmm_src);
+  void bit16_lut_compute_vector_fwd(const Xbyak::Zmm& zmm_src);
   void register_table_entries(const std::vector<postop_attr>& postop_attrs);
   void assert_check(const std::vector<postop_attr>& postop_attrs);
-  uint32_t get_int8_lut_term(int int8, const std::vector<postop_attr>& postop_attrs, data_type input_dt);
+  uint32_t get_bit8_lut_term(int integer, const std::vector<postop_attr>& postop_attrs, data_type output_dt);
+  uint32_t get_bit16_lut_term(int integer, const std::vector<postop_attr>& postop_attrs, data_type output_dt);
   std::string get_attr_idx_key(const postop_attr& attr);  // for get the key of alpha_idx,beta_idx,scale_idx map.
 
   void load_table_addr() { h->mov(p_table, l_table); }
@@ -133,10 +136,12 @@ class jit_eltwise_injector {
     tanh_saturation_lbound,
     tanh_pol_table,
     exchange_zmm_low256_high256,
-    int8_lut_term,
-    int8_64,
-    int8_255,
-    select_permt_idx,
+    bit8_lut_term,
+    bit8_64,
+    bit8_255,
+    bit16_lut_term,
+    bit16_32,
+    bit16_255,
     undef_key,
   };
 

nlp_toolkit/backends/neural_engine/SparseLib/include/jit_domain/jit_eltwiseop.hpp

@@ -66,23 +66,13 @@ class jit_eltwiseop_t : public jit_generator {
   Xbyak::Opmask remain_task_mask;
   Xbyak::Reg64 scratch_;
 
-  size_t dtype_size(postop_attr attr) {
-    // quantize only happend on first postop,we load the data from memory to zmm,in tail case,the offset is one byte;
-    // dequantize only happend on last postop,we store the data from zmm to memory,in taill case,the offset is also one
-    // byte.
-    if (attr.op_alg == postop_alg::int8_lut || attr.op_alg == postop_alg::quantize ||
-        attr.op_alg == postop_alg::dequantize)
-      return 1u;
-    switch (attr.dt) {
-      case data_type::fp32:
-        return 4u;
-      case data_type::bf16:
-        return 2u;
-    }
-  }
-
   size_t load_offset() {
-    if (param_.postop_attrs[0].op_alg == postop_alg::int8_lut) return 64u;  // special case:int8_lut
+    if (param_.postop_attrs[0].op_alg == postop_alg::eltop_int_lut && param_.postop_attrs[0].alpha == 8) {
+      return 64u;  // special case:bit8_lut
+    }
+    if (param_.postop_attrs[0].op_alg == postop_alg::eltop_int_lut && param_.postop_attrs[0].alpha == 16) {
+      return 32u;  // special case:bit16_lut
+    }
     auto head_dt = param_.postop_attrs.front().dt;
     switch (head_dt) {
       case data_type::fp32:
@@ -94,12 +84,11 @@ class jit_eltwiseop_t : public jit_generator {
   }
 
   size_t store_offset() {
-    // todo:the logic is little confuse,we need to optimize.
-    if (param_.postop_attrs.front().op_alg == postop_alg::int8_lut) return 64u;   // int8_lut case;
-    if (param_.postop_attrs.back().op_alg == postop_alg::quantize) return 16u;    // quantize case.
-    if (param_.postop_attrs.back().op_alg == postop_alg::dequantize) return 64u;  // dequantize case.
-    auto tail_dt = param_.postop_attrs.back().dt;
-    switch (tail_dt) {
+    // todo:except dequantize case, our zmm always full of result data and needs to be stored.
+    if (param_.postop_attrs.front().op_alg == postop_alg::eltop_int_lut) return 64u;  // lut case;
+    if (param_.postop_attrs.back().op_alg == postop_alg::quantize) return 16u;        // quantize case.
+    if (param_.postop_attrs.back().op_alg == postop_alg::dequantize) return 64u;      // dequantize case.
+    switch (param_.out_dt) {
       case data_type::fp32:
       case data_type::bf16:
         return 64u;
@@ -108,8 +97,9 @@ class jit_eltwiseop_t : public jit_generator {
 
   size_t process_element_num() {
     auto front_attr = param_.postop_attrs.front();
-    if (front_attr.op_alg == postop_alg::int8_lut) return 64;  // special case:int8_lut
-    switch (front_attr.dt) {
+    if (front_attr.op_alg == postop_alg::eltop_int_lut && front_attr.alpha == 8) return 64;   // special case:bit8_lut
+    if (front_attr.op_alg == postop_alg::eltop_int_lut && front_attr.alpha == 16) return 32;  // sepcial case:bit16_lut
+    switch (param_.in_dt) {
       case data_type::fp32:
         return 16;
       case data_type::bf16: