[WIP] x86/arm/mips/loongarch/riscv gemm block quantization optimization

src/layer/arm/gemm_arm.cpp

@@ -25,6 +25,10 @@ namespace ncnn {
 #endif
 #endif
 
+#if NCNN_WEIGHT_QUANT
+#include "gemm_wq_int8.h"
+#endif
+
 Gemm_arm::Gemm_arm()
 {
 #if __ARM_NEON
@@ -4582,10 +4586,160 @@ static int gemm_AT_BT_arm(const Mat& AT, const Mat& BT, const Mat& C, Mat& top_b
     return 0;
 }
 
+#if NCNN_WEIGHT_QUANT
+static int gemm_BT_arm_wq_int8(const Mat& A, const Mat& packed_B, const Mat& packed_B_descales, const Mat& input_scales, const Mat& C, Mat& top_blob, int broadcast_type_C, int N, int K, int block_size, int transA, int output_transpose, float alpha, float beta, int constant_TILE_M, int constant_TILE_N, int constant_TILE_K, int nT, const Option& opt)
+{
+    const int M = transA ? A.w : (A.dims == 3 ? A.c : A.h) * A.elempack;
+    const int block_count = (K + block_size - 1) / block_size;
+    const Mat BT = packed_B.reshape(K, N);
+    const Mat BT_descales = packed_B_descales.reshape(block_count, N);
+    int TILE_M, TILE_N, TILE_K;
+    get_optimal_tile_mnk_wq_int8(M, N, K, block_size, constant_TILE_M, constant_TILE_N, constant_TILE_K, TILE_M, TILE_N, TILE_K, nT);
+
+    const int mr = std::min(M, TILE_M);
+    const int nr = std::min(N, TILE_N);
+    const int nn_M = (M + TILE_M - 1) / TILE_M;
+    const int nn_N = (N + TILE_N - 1) / TILE_N;
+    const int nn_K = (K + TILE_K - 1) / TILE_K;
+    const float* input_scale_ptr = input_scales;
+    const size_t A_hstep = A.dims == 3 ? A.cstep : (size_t)A.w;
+
+    Mat topT(nr * mr, 1, nT, (size_t)4u, opt.workspace_allocator);
+    if (topT.empty())
+        return -100;
+
+    if (nT > nn_M)
+    {
+        Mat AT(K, mr, nn_M, (size_t)1u, opt.workspace_allocator);
+        Mat AT_descales(block_count, mr, nn_M, (size_t)4u, opt.workspace_allocator);
+        if (AT.empty() || AT_descales.empty())
+            return -100;
+
+        const int nn_MK = nn_M * nn_K;
+        #pragma omp parallel for num_threads(nT)
+        for (int ppik = 0; ppik < nn_MK; ppik++)
+        {
+            const int ppi = ppik / nn_K;
+            const int ppk = ppik % nn_K;
+            const int i = ppi * TILE_M;
+            const int k = ppk * TILE_K;
+            const int max_ii = std::min(M - i, TILE_M);
+            const int max_kk = std::min(K - k, TILE_K);
+            const int max_block_count = (max_kk + block_size - 1) / block_size;
+
+            Mat AT_channel = AT.channel(i / TILE_M);
+            Mat AT_descales_channel = AT_descales.channel(i / TILE_M);
+            Mat AT_tile(max_kk, max_ii, (signed char*)AT_channel + (size_t)k * mr, (size_t)1u);
+            Mat AT_descales_tile(max_block_count, max_ii, (float*)AT_descales_channel + (size_t)(k / block_size) * mr, (size_t)4u);
+
+            Mat A_tile = A;
+            A_tile.data = (unsigned char*)A_tile.data + (transA ? (size_t)k * A_hstep : (size_t)k) * sizeof(float);
+            const float* input_scale_tile = input_scale_ptr ? input_scale_ptr + k : 0;
+
+            if (transA)
+                transpose_quantize_A_tile_wq_int8(A_tile, AT_tile, AT_descales_tile, i, max_ii, max_kk, block_size, input_scale_tile);
+            else
+                quantize_A_tile_wq_int8(A_tile, AT_tile, AT_descales_tile, i, max_ii, max_kk, block_size, input_scale_tile);
+        }
+
+        const int nn_MN = nn_M * nn_N;
+        #pragma omp parallel for num_threads(nT)
+        for (int ppij = 0; ppij < nn_MN; ppij++)
+        {
+            const int ppi = ppij / nn_N;
+            const int ppj = ppij % nn_N;
+
+            const int i = ppi * TILE_M;
+            const int j = ppj * TILE_N;
+
+            const int max_ii = std::min(M - i, TILE_M);
+            const int max_jj = std::min(N - j, TILE_N);
+
+            Mat BT_tile = BT.row_range(j, max_jj);
+            Mat BT_descales_tile = BT_descales.row_range(j, max_jj);
+            Mat topT_tile = topT.channel(get_omp_thread_num());
+
+            Mat AT_channel = AT.channel(i / TILE_M);
+            Mat AT_descales_channel = AT_descales.channel(i / TILE_M);
+            for (int k = 0; k < K; k += TILE_K)
+            {
+                const int max_kk = std::min(K - k, TILE_K);
+                const int max_block_count = (max_kk + block_size - 1) / block_size;
+                Mat AT_tile(max_kk, max_ii, (signed char*)AT_channel + (size_t)k * mr, (size_t)1u);
+                Mat AT_descales_tile(max_block_count, max_ii, (float*)AT_descales_channel + (size_t)(k / block_size) * mr, (size_t)4u);
+
+                gemm_transB_packed_tile_wq_int8(AT_tile, AT_descales_tile, BT_tile, BT_descales_tile, topT_tile, max_ii, max_jj, K, k, max_kk, block_size);
+            }
+            if (output_transpose)
+                transpose_unpack_output_tile_wq_int8(topT_tile, C, top_blob, broadcast_type_C, i, max_ii, j, max_jj, N, alpha, beta);
+            else
+                unpack_output_tile_wq_int8(topT_tile, C, top_blob, broadcast_type_C, i, max_ii, j, max_jj, N, alpha, beta);
+        }
+    }
+    else
+    {
+        Mat ATX(K, mr, nT, (size_t)1u, opt.workspace_allocator);
+        Mat ATX_descales(block_count, mr, nT, (size_t)4u, opt.workspace_allocator);
+        if (ATX.empty() || ATX_descales.empty())
+            return -100;
+
+        #pragma omp parallel for num_threads(nT)
+        for (int ppi = 0; ppi < nn_M; ppi++)
+        {
+            const int i = ppi * TILE_M;
+            const int max_ii = std::min(M - i, TILE_M);
+
+            Mat AT_tile = ATX.channel(get_omp_thread_num());
+            Mat AT_descales_tile = ATX_descales.channel(get_omp_thread_num());
+            Mat topT_tile = topT.channel(get_omp_thread_num());
+
+            for (int j = 0; j < N; j += TILE_N)
+            {
+                const int max_jj = std::min(N - j, TILE_N);
+
+                Mat BT_tile = BT.row_range(j, max_jj);
+                Mat BT_descales_tile = BT_descales.row_range(j, max_jj);
+                for (int k = 0; k < K; k += TILE_K)
+                {
+                    const int max_kk = std::min(K - k, TILE_K);
+                    const int max_block_count = (max_kk + block_size - 1) / block_size;
+                    Mat AT_tile_k(max_kk, max_ii, (signed char*)AT_tile + (size_t)k * mr, (size_t)1u);
+                    Mat AT_descales_tile_k(max_block_count, max_ii, (float*)AT_descales_tile + (size_t)(k / block_size) * mr, (size_t)4u);
+
+                    if (j == 0)
+                    {
+                        Mat A_tile = A;
+                        A_tile.data = (unsigned char*)A_tile.data + (transA ? (size_t)k * A_hstep : (size_t)k) * sizeof(float);
+                        const float* input_scale_tile = input_scale_ptr ? input_scale_ptr + k : 0;
+
+                        if (transA)
+                            transpose_quantize_A_tile_wq_int8(A_tile, AT_tile_k, AT_descales_tile_k, i, max_ii, max_kk, block_size, input_scale_tile);
+                        else
+                            quantize_A_tile_wq_int8(A_tile, AT_tile_k, AT_descales_tile_k, i, max_ii, max_kk, block_size, input_scale_tile);
+                    }
+
+                    gemm_transB_packed_tile_wq_int8(AT_tile_k, AT_descales_tile_k, BT_tile, BT_descales_tile, topT_tile, max_ii, max_jj, K, k, max_kk, block_size);
+                }
+                if (output_transpose)
+                    transpose_unpack_output_tile_wq_int8(topT_tile, C, top_blob, broadcast_type_C, i, max_ii, j, max_jj, N, alpha, beta);
+                else
+                    unpack_output_tile_wq_int8(topT_tile, C, top_blob, broadcast_type_C, i, max_ii, j, max_jj, N, alpha, beta);
+            }
+        }
+    }
+
+    return 0;
+}
+#endif // NCNN_WEIGHT_QUANT
+
 int Gemm_arm::create_pipeline(const Option& opt)
 {
     if (weight_block_quantize)
     {
+#if NCNN_WEIGHT_QUANT
+        if (quantize_term / 100 == 8)
+            return create_pipeline_wq_int8(opt);
+#endif
         return 0;
     }
 
@@ -4745,10 +4899,150 @@ int Gemm_arm::create_pipeline(const Option& opt)
     return 0;
 }
 
+int Gemm_arm::destroy_pipeline(const Option& /*opt*/)
+{
+#if NCNN_WEIGHT_QUANT
+    BT_data_wq_int8.release();
+    BT_data_wq_int8_descales.release();
+#endif
+
+    return 0;
+}
+
+#if NCNN_WEIGHT_QUANT
+int Gemm_arm::create_pipeline_wq_int8(const Option& opt)
+{
+    if (!BT_data_wq_int8.empty())
+        return 0;
+
+    if (B_data.empty() || B_data_quantize_scales.empty())
+        return -100;
+
+    const int block_size_code = quantize_term % 10;
+    const int block_size = block_size_code == 0 ? 32 : block_size_code == 1 ? 64 : 128;
+
+    Mat BT_data_packed;
+    Mat BT_data_packed_descales;
+    int ret = pack_B_wq_int8(B_data, B_data_quantize_scales, BT_data_packed, BT_data_packed_descales, constantN, constantK, block_size, opt.num_threads);
+    if (ret != 0)
+        return ret;
+
+    BT_data_wq_int8 = BT_data_packed;
+    BT_data_wq_int8_descales = BT_data_packed_descales;
+
+    B_data.release();
+    B_data_quantize_scales.release();
+
+    return 0;
+}
+
+int Gemm_arm::forward_wq_int8(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_blobs, const Option& opt) const
+{
+    const Mat& A = bottom_blobs[0];
+    if (A.elemsize != 4u || A.elempack != 1)
+    {
+        NCNN_LOGE("Gemm unsupported input");
+        return -1;
+    }
+
+    if (transA && A.dims != 2)
+    {
+        NCNN_LOGE("Gemm unsupported input");
+        return -1;
+    }
+
+    const int K = transA ? A.h : A.w;
+    if (K != constantK)
+    {
+        NCNN_LOGE("Gemm weight block quantize K mismatch");
+        return -1;
+    }
+
+    const int M = transA ? A.w : A.dims == 3 ? A.c : A.h;
+    const int N = constantN;
+
+    Mat C;
+    int broadcast_type_C = -1;
+    if (constantC)
+    {
+        C = C_data;
+        broadcast_type_C = constant_broadcast_type_C;
+    }
+    else
+    {
+        if (bottom_blobs.size() == 2)
+            C = bottom_blobs[1];
+
+        if (!C.empty())
+        {
+            bool matched = false;
+            if (C.dims == 1 && C.w == 1)
+            {
+                broadcast_type_C = 0;
+                matched = true;
+            }
+            if (C.dims == 1 && C.w == M)
+            {
+                broadcast_type_C = 1;
+                matched = true;
+            }
+            if (C.dims == 1 && C.w == N)
+            {
+                broadcast_type_C = 4;
+                matched = true;
+            }
+            if (C.dims == 2 && C.w == 1 && C.h == M)
+            {
+                broadcast_type_C = 2;
+                matched = true;
+            }
+            if (C.dims == 2 && C.w == N && C.h == M)
+            {
+                broadcast_type_C = 3;
+                matched = true;
+            }
+            if (C.dims == 2 && C.w == N && C.h == 1)
+            {
+                broadcast_type_C = 4;
+                matched = true;
+            }
+
+            if (!matched || C.elemsize != 4u || C.elempack != 1)
+            {
+                NCNN_LOGE("Gemm unsupported C");
+                return -1;
+            }
+        }
+    }
+
+    if (!C.empty() && (C.elemsize != 4u || C.elempack != 1))
+    {
+        NCNN_LOGE("Gemm unsupported C");
+        return -1;
+    }
+
+    Mat& top_blob = top_blobs[0];
+    if (output_transpose)
+        top_blob.create(M, N, (size_t)4u, opt.blob_allocator);
+    else
+        top_blob.create(N, M, (size_t)4u, opt.blob_allocator);
+    if (top_blob.empty())
+        return -100;
+
+    const int block_size_code = quantize_term % 10;
+    const int block_size = block_size_code == 0 ? 32 : block_size_code == 1 ? 64 : 128;
+    return gemm_BT_arm_wq_int8(A, BT_data_wq_int8, BT_data_wq_int8_descales, B_data_input_scales, C, top_blob, broadcast_type_C, N, K, block_size, transA, output_transpose, alpha, beta, constant_TILE_M, constant_TILE_N, constant_TILE_K, opt.num_threads, opt);
+}
+#endif // NCNN_WEIGHT_QUANT
+
 int Gemm_arm::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_blobs, const Option& opt) const
 {
     if (weight_block_quantize)
     {
+#if NCNN_WEIGHT_QUANT
+        if (quantize_term / 100 == 8 && !BT_data_wq_int8.empty())
+            return forward_wq_int8(bottom_blobs, top_blobs, opt);
+#endif
         return Gemm::forward(bottom_blobs, top_blobs, opt);
     }
 

src/layer/arm/gemm_arm.h

@@ -14,6 +14,7 @@ class Gemm_arm : public Gemm
     Gemm_arm();
 
     virtual int create_pipeline(const Option& opt);
+    virtual int destroy_pipeline(const Option& opt);
 
     virtual int forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_blobs, const Option& opt) const;
 
@@ -34,6 +35,10 @@ class Gemm_arm : public Gemm
     int create_pipeline_int8(const Option& opt);
     int forward_int8(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_blobs, const Option& opt) const;
 #endif
+#if NCNN_WEIGHT_QUANT
+    int create_pipeline_wq_int8(const Option& opt);
+    int forward_wq_int8(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_blobs, const Option& opt) const;
+#endif
 
 public:
     int nT;
@@ -42,6 +47,11 @@ class Gemm_arm : public Gemm
     Mat CT_data;
 
     int input_elemtype; // 0=auto 1=fp32 2=fp16 3=bf16
+
+#if NCNN_WEIGHT_QUANT
+    Mat BT_data_wq_int8;
+    Mat BT_data_wq_int8_descales;
+#endif
 };
 
 } // namespace ncnn

src/layer/arm/gemm_arm_asimddp.cpp

@@ -14,6 +14,32 @@ namespace ncnn {
 #include "gemm_int8_bf16s.h"
 #endif
 
+#if NCNN_WEIGHT_QUANT
+#include "gemm_wq_int8.h"
+#endif
+
+#if NCNN_WEIGHT_QUANT
+int pack_B_wq_int8_asimddp(const Mat& B, const Mat& B_scales, Mat& BT, Mat& BT_descales, int N, int K, int block_size, int num_threads)
+{
+    return pack_B_wq_int8(B, B_scales, BT, BT_descales, N, K, block_size, num_threads);
+}
+
+void quantize_A_tile_wq_int8_asimddp(const Mat& A, Mat& AT_tile, Mat& AT_descales_tile, int i, int max_ii, int K, int block_size, const float* input_scale_ptr)
+{
+    quantize_A_tile_wq_int8(A, AT_tile, AT_descales_tile, i, max_ii, K, block_size, input_scale_ptr);
+}
+
+void transpose_quantize_A_tile_wq_int8_asimddp(const Mat& A, Mat& AT_tile, Mat& AT_descales_tile, int i, int max_ii, int K, int block_size, const float* input_scale_ptr)
+{
+    transpose_quantize_A_tile_wq_int8(A, AT_tile, AT_descales_tile, i, max_ii, K, block_size, input_scale_ptr);
+}
+
+void gemm_transB_packed_tile_wq_int8_asimddp(const Mat& AT_tile, const Mat& AT_descales_tile, const Mat& BT_tile, const Mat& BT_descales_tile, Mat& topT_tile, int max_ii, int max_jj, int K, int k, int max_kk, int block_size)
+{
+    gemm_transB_packed_tile_wq_int8(AT_tile, AT_descales_tile, BT_tile, BT_descales_tile, topT_tile, max_ii, max_jj, K, k, max_kk, block_size);
+}
+#endif // NCNN_WEIGHT_QUANT
+
 void pack_A_tile_int8_asimddp(const Mat& A, Mat& AT, int i, int max_ii, int k, int max_kk)
 {
     pack_A_tile_int8(A, AT, i, max_ii, k, max_kk);