Sparse VNNI tile fold (#223)

Author: yi1ding

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

@@ -75,6 +75,7 @@ class jit_spmm_vnni_t : public jit_generator {
   void store_intermediate_dst(dim_t m_start);
   void gen_subfunc_tile_prod();
   void gen_subfunc_dense_and_prod();
+  void load_dense_sparse_prod();
   void gen_subfunc_load_and_prod();
   void gen_subfunc_dst_epilogue();
   void handle_postop_escape_vmms();
@@ -118,6 +119,7 @@ class jit_spmm_vnni_t : public jit_generator {
   const Xbyak::Reg64& reg_scale = rbx;  // the scale
   const Xbyak::Opmask& reg_k1 = k1;
 
+  const Xbyak::Reg64& reg_k_ptr = param1;
   const Xbyak::Reg64& reg_tmp = r9;
   const Xbyak::Reg64& reg_dst_idx = r8;
   const Xbyak::Reg64& reg_m_idx = reg_tmp;

nlp_toolkit/backends/neural_engine/SparseLib/include/kernels/spmm_types.hpp

@@ -52,7 +52,8 @@ enum class subfunc_level : uint8_t {
   prod,            // use sub-function for tile product
   dense_and_prod,  // use fused sub-function for dense loading & tile product
   load_and_prod,   // use fused sub-function for dense loading & sparse loading & tile product
-  subfunc_level_MAX = load_and_prod
+  k_dims,         // a whole THxKxTW tile generates a constent size of code
+  subfunc_level_MAX = k_dims
 };
 
 /**

nlp_toolkit/backends/neural_engine/SparseLib/src/jit_domain/jit_spmm_vnni.cpp

@@ -155,44 +155,66 @@ void jit_spmm_vnni_t::repeat_THx4xTW_matmal(dim_t m_start) {
       mov(reg_seq_indices, reinterpret_cast<uint64_t>(dense_load_offsets.data() + indptr_lo - indptr_kernel_start));
       break;
     case ssd::subfunc_level::load_and_prod:
+    case ssd::subfunc_level::k_dims:
       mov(reg_seq_indices, reinterpret_cast<uint64_t>(dense_load_offsets.data() + indptr_lo - indptr_kernel_start));
       mov(reg_wei, reinterpret_cast<uint64_t>(param_.weight + param_.blocksize[0] * param_.blocksize[1] * indptr_lo));
       break;
     default:
       break;
   }
-
-  // kp (k-idx pointer is the idx of nnz blocks of the current row)
-  for (int64_t kp_lo = 0; kp_lo < nnz; kp_lo += spns::ADJ) {
-    const int64_t kp_hi = std::min(kp_lo + spns::ADJ, nnz);  // end of k-index pointer (noninclusive)
-    dim_t element_offset = param_.blocksize[0] * param_.blocksize[1] * indptr_lo + kp_lo * TH();
-
-    // Step 1: load dense (activation). Note that k_indices length is processed.00-00
-    // Step 2: load sparse (weight) and reorder data for that.
-    // Step 3: tile product. Note that k_indices length is processed.
-    // A tile product can calculate at least 1 row and 16 columns of DST.
-    // Min tile calculation: Tile width/height is 1, compute (1, ADJ) x (ADJ, 16) = (1, 16) matmul.
-    switch (param_.sub_func) {
-      case ssd::subfunc_level::none:
-        load_dense({k_indices.begin() + kp_lo, k_indices.begin() + kp_hi});
-        load_sparse(reg_wei, element_offset * sizeof(decltype(*param_.weight)));
-        tile_product(TH(), TW());
-        break;
-      case ssd::subfunc_level::prod:
-        load_dense({k_indices.begin() + kp_lo, k_indices.begin() + kp_hi});
-        load_sparse(reg_wei, element_offset * sizeof(decltype(*param_.weight)));
-        call(sfptr_tile_prod_);
-        break;
-      case ssd::subfunc_level::dense_and_prod:
-        load_sparse(reg_wei, element_offset * sizeof(decltype(*param_.weight)));
-        call(sfptr_dense_and_prod_);
-        break;
-      case ssd::subfunc_level::load_and_prod:
-        call(sfptr_load_and_prod_);
-        break;
-      default:
+  switch (param_.sub_func) {
+    case ssd::subfunc_level::none:
+    case ssd::subfunc_level::prod:
+    case ssd::subfunc_level::dense_and_prod:
+    case ssd::subfunc_level::load_and_prod:
+      // kp (k-idx pointer is the idx of nnz blocks of the current row)
+      for (int64_t kp_lo = 0; kp_lo < nnz; kp_lo += spns::ADJ) {
+        const int64_t kp_hi = std::min(kp_lo + spns::ADJ, nnz);  // end of k-index pointer (noninclusive)
+        dim_t element_offset = param_.blocksize[0] * param_.blocksize[1] * indptr_lo + kp_lo * TH();
+
+        // Step 1: load dense (activation). Note that k_indices length is processed.00-00
+        // Step 2: load sparse (weight) and reorder data for that.
+        // Step 3: tile product. Note that k_indices length is processed.
+        // A tile product can calculate at least 1 row and 16 columns of DST.
+        // Min tile calculation: Tile width/height is 1, compute (1, ADJ) x (ADJ, 16) = (1, 16) matmul.
+        switch (param_.sub_func) {
+          case ssd::subfunc_level::none:
+            load_dense({k_indices.begin() + kp_lo, k_indices.begin() + kp_hi});
+            load_sparse(reg_wei, element_offset * sizeof(decltype(*param_.weight)));
+            tile_product(TH(), TW());
+            break;
+          case ssd::subfunc_level::prod:
+            load_dense({k_indices.begin() + kp_lo, k_indices.begin() + kp_hi});
+            load_sparse(reg_wei, element_offset * sizeof(decltype(*param_.weight)));
+            call(sfptr_tile_prod_);
+            break;
+          case ssd::subfunc_level::dense_and_prod:
+            load_sparse(reg_wei, element_offset * sizeof(decltype(*param_.weight)));
+            call(sfptr_dense_and_prod_);
+            break;
+          case ssd::subfunc_level::load_and_prod:
+            call(sfptr_load_and_prod_);
+            break;
+          default:
+            break;
+        }
+      }
+      break;
+    case ssd::subfunc_level::k_dims:
+      if (nnz > 0) {  // at least one iteration
+        xor_(reg_k_ptr, reg_k_ptr);
+        add(reg_dense, reg_n_idx);  // reg_dense += reg_n_idx * BYTE1
+
+        Xbyak::Label L_adj_k_loop;
+        L(L_adj_k_loop);
+        load_dense_sparse_prod();
+        add(reg_k_ptr, spns::ADJ);
+        cmp(reg_k_ptr, static_cast<int>(nnz));
+        jl(L_adj_k_loop);  // Loop-N2 end.
+
+        sub(reg_dense, reg_n_idx);  // reg_dense = reg_n_idx * BYTE1
         break;
-    }
+      }
   }
 }
 
@@ -316,10 +338,13 @@ void jit_spmm_vnni_t::gen_subfunc_dense_and_prod() {
   ret();
 }
 
-void jit_spmm_vnni_t::gen_subfunc_load_and_prod() {
-  sfptr_load_and_prod_ = getCurr();
-  add(reg_dense, reg_n_idx);  // reg_dense += reg_n_idx * BYTE1
-
+/**
+ * Required registers:
+ *  reg_dense - the start of the current row of dense matrix
+ *  reg_seq_indices - the start of offset for each TW, it will be updated after read
+ *  reg_wei - the start of weight matrix, it will be updated after read
+ */
+void jit_spmm_vnni_t::load_dense_sparse_prod() {
   constexpr size_t idx_size = sizeof(decltype(param_.indices)::value_type);
   mov(reg_addr_tmp[0], qword[reg_seq_indices + 0 * idx_size]);
   mov(reg_addr_tmp[1], qword[reg_seq_indices + 1 * idx_size]);
@@ -353,8 +378,16 @@ void jit_spmm_vnni_t::gen_subfunc_load_and_prod() {
       // tile prod
       vpdpbusd(dst_tile_Vmm(i, j), TW_Vmm(j), TH_Vmm(i));
     }
+    // update reg_wei in the middle
     if (j == TW() / 2) add(reg_wei, TH() * spns::ADJ * wei_size);
   }
+}
+
+void jit_spmm_vnni_t::gen_subfunc_load_and_prod() {
+  sfptr_load_and_prod_ = getCurr();
+  add(reg_dense, reg_n_idx);  // reg_dense += reg_n_idx * BYTE1
+
+  load_dense_sparse_prod();
 
   sub(reg_dense, reg_n_idx);  // reg_dense = reg_n_idx * BYTE1
   ret();
@@ -368,6 +401,7 @@ void jit_spmm_vnni_t::generate() {
   gen_subfunc_dst_epilogue();
   switch (param_.sub_func) {
     case ssd::subfunc_level::none:
+    case ssd::subfunc_level::k_dims:
       break;
     case ssd::subfunc_level::prod:
       gen_subfunc_tile_prod();

nlp_toolkit/backends/neural_engine/test/SparseLib/benchmark/ci/inputs/ci_vnni_input

@@ -1,36 +1,38 @@
 ### ncores_per_inst operator output_channel input_channel bsxseq sparsity micro_bs is_fp32_out has_append_sum micro_oc sub_func_level postop
 
 # bert mini sp90 bs=1 c4
-4 vnni 256 256 128 0.9 -1 0 0 -1 0
-4 vnni 256 1024 128 0.9 -1 0 0 -1 0
-4 vnni 1024 256 128 0.9 -1 0 0 -1 0
+4 vnni 256 256 128 0.9 -1 0 0 -1 -1
+4 vnni 256 1024 128 0.9 -1 0 0 -1 -1
+4 vnni 1024 256 128 0.9 -1 0 0 -1 -1
 
 # bert large sp90 bs=1 c4
-4 vnni 1024 1024 128 0.9 -1 0 0 -1 0
-4 vnni 1024 4096 128 0.9 -1 0 0 -1 0
-4 vnni 4096 1024 128 0.9 -1 0 0 -1 0
+4 vnni 1024 1024 128 0.9 -1 0 0 -1 -1
+4 vnni 1024 4096 128 0.9 -1 0 0 -1 -1
+4 vnni 4096 1024 128 0.9 -1 0 0 -1 -1
 
 # distilbert sp80 bs=8 c4/7
-4 vnni 768 768 256 0.8 -1 0 0 -1 0
-4 vnni 768 3072 256 0.8 -1 0 0 -1 0
-4 vnni 3072 768 256 0.8 -1 0 0 -1 0
-7 vnni 768 768 256 0.8 -1 0 0 -1 0
-7 vnni 768 3072 256 0.8 -1 0 0 -1 0
-7 vnni 3072 768 256 0.8 -1 0 0 -1 0
-7 vnni 768 768 256 0.8 -1 0 0 -1 0 gelu
-7 vnni 768 3072 256 0.8 -1 0 0 -1 0 gelu
-7 vnni 3072 768 256 0.8 -1 0 0 -1 0 gelu
+4 vnni 768 768 256 0.8 -1 0 0 -1 -1
+4 vnni 768 3072 256 0.8 -1 0 0 -1 -1
+4 vnni 3072 768 256 0.8 -1 0 0 -1 -1
+7 vnni 768 768 256 0.8 -1 0 0 -1 -1
+7 vnni 768 3072 256 0.8 -1 0 0 -1 -1
+7 vnni 3072 768 256 0.8 -1 0 0 -1 -1
+7 vnni 768 768 256 0.8 -1 0 0 -1 -1 gelu
+7 vnni 768 3072 256 0.8 -1 0 0 -1 -1 gelu
+7 vnni 3072 768 256 0.8 -1 0 0 -1 -1 gelu
 
 # distilbert sp80 bs=8 c7 micro_bs/micro_oc
-4 vnni 768 768 256 0.8 128 0 0 -1 0
-4 vnni 3072 768 256 0.8 -1 0 0 128 0
+4 vnni 768 768 256 0.8 128 0 0 -1 -1
+4 vnni 3072 768 256 0.8 -1 0 0 128 -1
 
 # cases for differnt op_attrs
+8 vnni 3072 768 256 0.8 -1 0 0 -1 -1
+8 vnni 3072 768 256 0.8 128 0 0 -1 -1
+8 vnni 3072 768 256 0.8 -1 1 0 -1 -1
+8 vnni 3072 768 256 0.8 -1 1 1 -1 -1
+8 vnni 3072 768 256 0.8 -1 0 0 128 -1
 8 vnni 3072 768 256 0.8 -1 0 0 -1 0
-8 vnni 3072 768 256 0.8 128 0 0 -1 0
-8 vnni 3072 768 256 0.8 -1 1 0 -1 0
-8 vnni 3072 768 256 0.8 -1 1 1 -1 0
-8 vnni 3072 768 256 0.8 -1 0 0 128 0
 8 vnni 3072 768 256 0.8 -1 0 0 -1 1
 8 vnni 3072 768 256 0.8 -1 0 0 -1 2
 8 vnni 3072 768 256 0.8 -1 0 0 -1 3
+8 vnni 3072 768 256 0.8 -1 0 0 -1 4

nlp_toolkit/backends/neural_engine/test/gtest/SparseLib/test_spmm_vnni_kernel.cpp

@@ -397,6 +397,7 @@ static auto case_func = []() {
       cases.push_back({gen_case(32, 32, 128, .7f, -1, nthr, dt::fp32, {{"sub_func", "1"}})});
       cases.push_back({gen_case(32, 32, 128, .7f, -1, nthr, dt::fp32, {{"sub_func", "2"}})});
       cases.push_back({gen_case(32, 32, 128, .7f, -1, nthr, dt::fp32, {{"sub_func", "3"}})});
+      cases.push_back({gen_case(32, 32, 128, .7f, -1, nthr, dt::fp32, {{"sub_func", "4"}})});
 
       // case: sparse: s8xu8+s32=s8, weight(M, K) * activation(K, N) + bias(M, 1) = dst(M, N)
       cases.push_back({gen_case(32, 32, 128, .7f, -1, nthr, dt::s8)});