feature: add support of RISCV64_SPACEMIT_IME2

Change-Id: I07c3e0dbb9bc10a11bcb92df1bbad75077c0e06a

Author: ZhaoLikunSpacmit

ggml/src/ggml-cpu/CMakeLists.txt

@@ -21,14 +21,14 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
             ggml-cpu/ggml-cpu-traits.h
             ggml-cpu/ggml-cpu-impl.h
         )
-        
+
         if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "riscv64")
             include(FetchContent)
             # TODO replace with git repo
             FetchContent_Declare(
                 onnxruntime
                 GIT_REPOSITORY ssh://$ENV{GERRIT_USER}@gerrit.dc.com:29418/DSA/onnxruntime
-                GIT_TAG "c17089e2e45067e24911d95611d2196a3dd63694"
+                GIT_TAG "7935d26a2ef0afa307e39b4c8a2ed438d281e5bd"
             )
             # FetchContent_Declare(
             #     onnxruntime
@@ -348,6 +348,7 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
         message(STATUS "RISC-V detected")
         if (GGML_RVV)
             list(APPEND ARCH_FLAGS -march=rv64gcv -mabi=lp64d)
+            list(APPEND ARCH_DEFINITIONS RISCV64_SPACEMIT_IME2)
         endif()
     else()
         message(STATUS "Unknown architecture")

ggml/src/ggml-cpu/ggml-cpu-riscv64-spacemit.cpp

@@ -80,7 +80,11 @@ void SQ4BitGemm_CompInt8(
 
   size_t CountN;
 
+#if defined(RISCV64_SPACEMIT_IME1)
   const size_t ComputeBlockCountN = RangeCountM == 1 ? RangeCountN : 16;
+#elif defined(RISCV64_SPACEMIT_IME2)
+  const size_t ComputeBlockCountN = RangeCountM == 1 ? RangeCountN : 32;
+#endif
 
   for (size_t n = 0; n < RangeCountN; n += CountN) {
     CountN = std::min(RangeCountN - n, ComputeBlockCountN);
@@ -279,6 +283,8 @@ struct block {
 };
 
 // control size
+static_assert(sizeof(block<4, 32>) == 32 * sizeof(ggml_half) + QK4_0 * 16, "wrong block<4,32> size/padding");
+static_assert(sizeof(block<8, 32>) == 32 * sizeof(ggml_half) + QK4_0 * 32, "wrong block<8,32> size/padding");
 static_assert(sizeof(block<4, 16>) == 16 * sizeof(ggml_half) + QK4_0 * 8, "wrong block<4,16> size/padding");
 static_assert(sizeof(block<8, 16>) == 16 * sizeof(ggml_half) + QK4_0 * 16, "wrong block<8,16> size/padding");
 
@@ -296,24 +302,55 @@ static block_q4_0x16 make_block_q4_0x16(block_q4_0* in, unsigned int blck_size_i
   for (int i = 0; i < 16; i++) {
     // [0, 15], in.d & 0x0F
     for (int j = 0; j < QK4_0 / 4; j++) {
-      // [b0 b16] ......... [b8 b24] ......... [b15 b31]
-      // [b0 b8] ......... [b7 b15]
+      //src [b0 b16] ......... [b8 b24] ......... [b15 b31]
+      //dst [b0 b8] ......... [b7 b15]
       out.qs[i * QK4_0 / 4 + j] = (in[i].qs[j] & 0x0F) | ((in[i].qs[j + QK4_0 / 4] & 0x0F) << 4);
     }
   }
 
   for (int i = 0; i < 16; i++) {
     // [16, 31], in.d & 0xF0
     for (int j = 0; j < QK4_0 / 4; j++) {
-      // [b0 b16] ......... [b8 b24] ......... [b15 b31]
-      // [b16 b24] ......... [b23 b31]
+      //src [b0 b16] ......... [b8 b24] ......... [b15 b31]
+      //dst [b16 b24] ......... [b23 b31]
       out.qs[4 * QK4_0 + i * QK4_0 / 4 + j] = ((in[i].qs[j] & 0xF0) >> 4) | (in[i].qs[j + QK4_0 / 4] & 0xF0);
     }
   }
 
   return out;
 }
 
+using block_q4_0x32 = block<4, 32>;
+using block_q8_0x32 = block<8, 32>;
+static block_q4_0x32 make_block_q4_0x32(block_q4_0* in, unsigned int blck_size_interleave) {
+  block_q4_0x32 out;
+  assert(QK4_0 / blck_size_interleave == 1);
+
+  for (int i = 0; i < 32; i++) { //  zhaolikun [check]
+    out.d[i] = in[i].d;
+  }
+
+  for (int i = 0; i < 32; i++) {
+    // [0, 15], in.d & 0x0F
+    for (int j = 0; j < QK4_0/4; j++) {
+      //src [b0 b16] ......... [b8 b24] ......... [b15 b31]
+      //dst [b0 b1] .........  [b14 b15]
+      out.qs[i * QK4_0/2 + j] = (in[i].qs[j*2] & 0x0F) | ((in[i].qs[j*2 + 1] & 0x0F) << 4);
+    }
+  }
+
+  for (int i = 0; i < 32; i++) {
+    // [16, 31], in.d & 0xF0
+    for (int j = 0; j < QK4_0/4; j++) {
+      //src [b0 b16] ......... [b8 b24] ......... [b15 b31]
+      //dst [b16 b17] ......... [b30 b31]
+      out.qs[i * QK4_0/2 + QK4_0/4 + j] = ((in[i].qs[j*2] & 0xF0)>>4) | (in[i].qs[j*2 + 1] & 0xF0);
+    }
+  }
+
+  return out;
+}
+
 static int repack_q4_0_to_q4_0_16_bl(struct ggml_tensor* t, int interleave_block, const void* GGML_RESTRICT data, size_t data_size) {
   GGML_ASSERT(t->type == GGML_TYPE_Q4_0);
   GGML_ASSERT(interleave_block == 16);
@@ -346,6 +383,38 @@ static int repack_q4_0_to_q4_0_16_bl(struct ggml_tensor* t, int interleave_block
   GGML_UNUSED(data_size);
 }
 
+static int repack_q4_0_to_q4_0_32_bl(struct ggml_tensor* t, int interleave_block, const void* GGML_RESTRICT data, size_t data_size) {
+  GGML_ASSERT(t->type == GGML_TYPE_Q4_0);
+  GGML_ASSERT(interleave_block == 32); // unused
+
+  constexpr int nrows_interleaved = 32;
+
+  block_q4_0x32* dst = (block_q4_0x32*)t->data;
+  const block_q4_0* src = (const block_q4_0*)data;
+  block_q4_0 dst_tmp[32];
+  int nrow = ggml_nrows(t);
+  int nblocks = t->ne[0] / QK4_0;
+
+  GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_0));
+
+  if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % QK4_0 != 0) {
+    return -1;
+  }
+
+  for (int b = 0; b < nrow; b += nrows_interleaved) {
+    for (int64_t x = 0; x < nblocks; x++) {
+      for (int i = 0; i < nrows_interleaved; i++) {
+        dst_tmp[i] = src[x + i * nblocks];
+      }
+      *dst++ = make_block_q4_0x32(dst_tmp, interleave_block);
+    }
+    src += nrows_interleaved * nblocks;
+  }
+  return 0;
+
+  GGML_UNUSED(data_size);
+}
+
 namespace ggml::cpu::riscv64_spacemit {
 
 template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS>
@@ -355,6 +424,10 @@ template <>
 int repack<block_q4_0, 8, 16>(struct ggml_tensor* t, const void* data, size_t data_size) {
   return repack_q4_0_to_q4_0_16_bl(t, 16, data, data_size);
 }
+template <>
+int repack<block_q4_0, 16, 32>(struct ggml_tensor* t, const void* data, size_t data_size) {
+  return repack_q4_0_to_q4_0_32_bl(t, 32, data, data_size);
+}
 
 class tensor_traits_base : public ggml::cpu::tensor_traits {
  public:
@@ -707,15 +780,22 @@ class tensor_traits_common : public tensor_traits_base {
 };
 
 static const tensor_traits<block_q4_0, 8, 16> q4_0_16x8_q8_0;
+static const tensor_traits<block_q4_0, 16, 32> q4_0_32x16_q8_0;
 static const tensor_traits_common rvv_impl;
 
 }  // namespace ggml::cpu::riscv64_spacemit
 
 static const ggml::cpu::tensor_traits* ggml_riscv64_spacemit_get_optimal_repack_type(const struct ggml_tensor* cur) {
   if (cur->type == GGML_TYPE_Q4_0) {
-    if (cur->ne[1] % 16 == 0) {
-      return &ggml::cpu::riscv64_spacemit::q4_0_16x8_q8_0;
-    }
+    #if defined(RISCV64_SPACEMIT_IME1)
+        if (cur->ne[1] % 16 == 0) {
+          return &ggml::cpu::riscv64_spacemit::q4_0_16x8_q8_0;
+        }
+    #elif defined(RISCV64_SPACEMIT_IME2)
+        if (cur->ne[1] % 32 == 0) {
+          return &ggml::cpu::riscv64_spacemit::q4_0_32x16_q8_0;
+        }
+    #endif
   } else if (cur->type == GGML_TYPE_F32) {
     return &ggml::cpu::riscv64_spacemit::rvv_impl;
   }

ggml/src/ggml-cpu/onnxruntime_mlas/CMakeLists.txt

@@ -59,7 +59,7 @@ elseif(NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten")
       set(ARM64 TRUE)
     elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^riscv64.*")
       set(RISCV64 TRUE)
-      set(RISCV64_SPACEMIT_IME_SPEC RISCV64_SPACEMIT_IME1)
+      set(RISCV64_SPACEMIT_IME_SPEC RISCV64_SPACEMIT_IME2)
     elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(i.86|x86?)$")
       set(X86 TRUE)
     elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|amd64)$")
@@ -768,7 +768,7 @@ endif()
           ${MLAS_SRC_DIR}/layernorm_rvv.cpp
         )
 
-        if (RISCV64_SPACEMIT_IME_SPEC)
+        if (RISCV64_SPACEMIT_IME_SPEC STREQUAL "RISCV64_SPACEMIT_IME1")
           target_compile_definitions(onnxruntime_mlas PRIVATE ${RISCV64_SPACEMIT_IME_SPEC})
           set(mlas_platform_srcs
             ${mlas_platform_srcs}
@@ -779,6 +779,17 @@ endif()
           )
         endif()
 
+        if (RISCV64_SPACEMIT_IME_SPEC STREQUAL "RISCV64_SPACEMIT_IME2")
+          target_compile_definitions(onnxruntime_mlas PRIVATE ${RISCV64_SPACEMIT_IME_SPEC})
+          set(mlas_platform_srcs
+            ${mlas_platform_srcs}
+            ${MLAS_SRC_DIR}/qgemm_kernel_spacemit_ime2.cpp
+            ${MLAS_SRC_DIR}/sqnbitgemm_kernel_spacemit_ime2.cpp
+            ${MLAS_SRC_DIR}/sqnbitgemm_kernel_spacemit_ime2_int8.cpp
+            ${MLAS_SRC_DIR}/sqnbitgemm_kernel_spacemit_ime_fp32.cpp
+          )
+        endif()
+
         if(NOT ONNXRUNTIME_MLAS_MULTI_ARCH)
           set(MLAS_SOURCE_IS_NOT_SET 0)
         endif()