[FORK][FEATURE] DQ IP: precompute grouped src sums

Author: dmitry-gorokhov

src/common/memory_tracking.hpp

@@ -305,6 +305,7 @@ enum {
     key_decompression_zero_points,
     key_src_quantized,
     key_src_dequantized_scales,
+    key_src_grouped_sum,
     // These two keys should always be the last ones,
     // even though they are not in alphabetical order
     key_nested,

src/cpu/x64/brgemm/brgemm.cpp

@@ -82,7 +82,8 @@ void brgemm_desc_t::cleanup_dst_md() {
 void brgemm_kernel_execute(const brgemm_kernel_t *brg_kernel, int bs,
         const brgemm_batch_element_t *batch, void *ptr_C, void *scratch,
         const brgemm_dynamic_values_t *dynamic_values,
-        const void *ptr_wei_scales, const void *ptr_wei_zero_points, const void *ptr_src_scales, size_t ic) {
+        const void *ptr_wei_scales, const void *ptr_wei_zero_points,
+        const void *ptr_src_scales, const void *ptr_src_grouped_sum, size_t ic) {
     brgemm_kernel_params_t brgemm_p;
 
     brgemm_p.batch = batch;
@@ -105,6 +106,7 @@ void brgemm_kernel_execute(const brgemm_kernel_t *brg_kernel, int bs,
     brgemm_p.ptr_wei_scales = ptr_wei_scales;
     brgemm_p.ptr_wei_zero_points = ptr_wei_zero_points;
     brgemm_p.ptr_src_scales = ptr_src_scales;
+    brgemm_p.ptr_src_grouped_sum = ptr_src_grouped_sum;
     brgemm_p.ic = ic;
 
     assert(brg_kernel);
@@ -116,7 +118,8 @@ void brgemm_kernel_execute(const brgemm_kernel_t *brg_kernel, int bs,
         const void *addr_A, const void *addr_B,
         const brgemm_batch_element_t *batch, void *ptr_C, void *scratch,
         const brgemm_dynamic_values_t *dynamic_values,
-        const void *ptr_wei_scales, const void *ptr_wei_zero_points, const void *ptr_src_scales, size_t ic) {
+        const void *ptr_wei_scales, const void *ptr_wei_zero_points,
+        const void *ptr_src_scales, const void *ptr_src_grouped_sum, size_t ic) {
     brgemm_kernel_params_t brgemm_p;
 
     brgemm_p.batch = batch;
@@ -133,6 +136,7 @@ void brgemm_kernel_execute(const brgemm_kernel_t *brg_kernel, int bs,
     brgemm_p.ptr_wei_scales = ptr_wei_scales;
     brgemm_p.ptr_wei_zero_points = ptr_wei_zero_points;
     brgemm_p.ptr_src_scales = ptr_src_scales;
+    brgemm_p.ptr_src_grouped_sum = ptr_src_grouped_sum;
     brgemm_p.ic = ic;
     if (dynamic_values) {
         brgemm_p.dynamic_LDA = dynamic_values->dynamic_LDA;
@@ -148,7 +152,8 @@ void brgemm_kernel_execute_postops(const brgemm_kernel_t *brg_kernel, int bs,
         const brgemm_batch_element_t *batch, void *ptr_C, void *ptr_D,
         const brgemm_post_ops_data_t &post_ops_data, void *scratch,
         const brgemm_dynamic_values_t *dynamic_values,
-        const void *ptr_wei_scales, const void *ptr_wei_zero_points, const void *ptr_src_scales, size_t ic) {
+        const void *ptr_wei_scales, const void *ptr_wei_zero_points,
+        const void *ptr_src_scales, const void *ptr_src_grouped_sum, size_t ic) {
     brgemm_kernel_params_t brgemm_p;
 
     brgemm_p.batch = batch;
@@ -178,6 +183,7 @@ void brgemm_kernel_execute_postops(const brgemm_kernel_t *brg_kernel, int bs,
     brgemm_p.ptr_wei_scales = ptr_wei_scales;
     brgemm_p.ptr_wei_zero_points = ptr_wei_zero_points;
     brgemm_p.ptr_src_scales = ptr_src_scales;
+    brgemm_p.ptr_src_grouped_sum = ptr_src_grouped_sum;
     brgemm_p.ic = ic;
     if (dynamic_values) {
         brgemm_p.dynamic_LDA = dynamic_values->dynamic_LDA;
@@ -194,7 +200,8 @@ void brgemm_kernel_execute_postops(const brgemm_kernel_t *brg_kernel, int bs,
         const brgemm_batch_element_t *batch, void *ptr_C, void *ptr_D,
         const brgemm_post_ops_data_t &post_ops_data, void *scratch,
         const brgemm_dynamic_values_t *dynamic_values,
-        const void *ptr_wei_scales, const void *ptr_wei_zero_points, const void *ptr_src_scales, size_t ic) {
+        const void *ptr_wei_scales, const void *ptr_wei_zero_points,
+        const void *ptr_src_scales, const void *ptr_src_grouped_sum, size_t ic) {
     brgemm_kernel_params_t brgemm_p;
 
     brgemm_p.batch = batch;
@@ -224,6 +231,7 @@ void brgemm_kernel_execute_postops(const brgemm_kernel_t *brg_kernel, int bs,
     brgemm_p.ptr_wei_scales = ptr_wei_scales;
     brgemm_p.ptr_wei_zero_points = ptr_wei_zero_points;
     brgemm_p.ptr_src_scales = ptr_src_scales;
+    brgemm_p.ptr_src_grouped_sum = ptr_src_grouped_sum;
     brgemm_p.ic = ic;
     if (dynamic_values) {
         brgemm_p.dynamic_LDA = dynamic_values->dynamic_LDA;
@@ -318,6 +326,12 @@ status_t brgemm_desc_init(brgemm_desc_t *brg, cpu_isa_t isa,
 
     CHECK(brgemm_blocking(brg));
 
+    brg->src_sum_group_size = wei_d.dims()[1];
+    if (brg->with_src_dyn_quant) {
+        brg->src_sum_group_size = brg->rd_block;
+        brg->src_grouped_sum_stride = div_up(wei_d.dims()[1], brg->src_sum_group_size);
+    }
+
     // avx2_vnni_2 kernel with xf16 data type requires blocked weights.
     if (brg->isa_impl == avx2_vnni_2 && brg->is_xf16()
             && brg->LDB % brg->ld_block > 0)

src/cpu/x64/brgemm/brgemm.hpp

@@ -175,7 +175,7 @@ void DNNL_API brgemm_kernel_execute(const brgemm_kernel_t *brg_kernel, int bs,
         void *scratch = nullptr,
         const brgemm_dynamic_values_t *dynamic_values = nullptr,
         const void *ptr_wei_scales = nullptr, const void *ptr_wei_zero_points = nullptr,
-        const void *ptr_src_scales = nullptr, size_t ic = 0);
+        const void *ptr_src_scales = nullptr, const void *ptr_src_grouped_sum = nullptr, size_t ic = 0);
 
 /// Execute BRGEMM kernel (brgemm_offs and brgemm_strd version)
 ///
@@ -205,7 +205,7 @@ void brgemm_kernel_execute(const brgemm_kernel_t *brg_kernel, int bs,
         void *scratch = nullptr,
         const brgemm_dynamic_values_t *dynamic_values = nullptr,
         const void *ptr_wei_scales = nullptr, const void *ptr_wei_zero_points = nullptr,
-        const void *ptr_src_scales = nullptr, size_t ic = 0);
+        const void *ptr_src_scales = nullptr, const void *ptr_src_grouped_sum = nullptr, size_t ic = 0);
 
 /// Execute BRGEMM kernel (brgemm_addr version)
 ///
@@ -234,7 +234,7 @@ void DNNL_API brgemm_kernel_execute_postops(const brgemm_kernel_t *brg_kernel,
         const brgemm_post_ops_data_t &post_ops_data, void *scratch = nullptr,
         const brgemm_dynamic_values_t *dynamic_values = nullptr,
         const void *ptr_wei_scales = nullptr, const void *ptr_wei_zero_points = nullptr,
-        const void *ptr_src_scales = nullptr, size_t ic = 0);
+        const void *ptr_src_scales = nullptr, const void *ptr_src_grouped_sum = nullptr, size_t ic = 0);
 
 /// Execute BRGEMM kernel (brgemm_offs and brgemm_strd version)
 ///
@@ -267,7 +267,7 @@ void DNNL_API brgemm_kernel_execute_postops(const brgemm_kernel_t *brg_kernel,
         const brgemm_post_ops_data_t &post_ops_data, void *scratch = nullptr,
         const brgemm_dynamic_values_t *dynamic_values = nullptr,
         const void *ptr_wei_scales = nullptr, const void *ptr_wei_zero_points = nullptr,
-        const void *ptr_src_scales = nullptr, size_t ic = 0);
+        const void *ptr_src_scales = nullptr, const void *ptr_src_grouped_sum = nullptr, size_t ic = 0);
 
 /// AMX utilities: Creates a palette based on BRGEMM descriptor
 ///

src/cpu/x64/brgemm/brgemm_types.hpp

@@ -321,6 +321,8 @@ struct brgemm_desc_t {
     bool with_src_dyn_quant = false;
     int src_scales_group_size = 0;
     int src_scales_stride = 0;
+    int src_sum_group_size = 0;
+    int src_grouped_sum_stride = 0;
 
     bool is_row_major() const {
         assert(layout != brgemm_layout_undef);
@@ -500,6 +502,7 @@ struct brgemm_kernel_params_t {
     const void *ptr_wei_scales = nullptr;
     const void *ptr_wei_zero_points = nullptr;
     const void *ptr_src_scales = nullptr;
+    const void *ptr_src_grouped_sum = nullptr;
     size_t ic;
     dim_t dynamic_LDA = 0;
     dim_t dynamic_LDB = 0;

src/cpu/x64/brgemm/brgemm_utils.cpp

@@ -231,7 +231,7 @@ int calculate_max_bcast_block(brgemm_desc_t *brg, const int adj_ld_block2) {
     if (one_of(brg->dt_b, data_type::f4_e2m1) && brg->isa_impl == avx2) max_bcast_block -= 2;
     if (one_of(brg->dt_b, data_type::nf4, data_type::f4_e2m1) && brg->isa_impl != avx2) max_bcast_block -= 1;
     if (brg->with_wei_decomp_zero_points && brg->wei_decomp_zero_points_stride == 0) max_bcast_block -= 1;
-    if (brg->with_src_dyn_quant) max_bcast_block -= 2;
+    if (brg->with_src_dyn_quant) max_bcast_block -= 1;
     if (brg->with_src_dyn_quant && brg->with_wei_decomp_zero_points && brg->wei_decomp_zero_points_stride != 0) max_bcast_block -= adj_ld_block2;
 
     max_bcast_block /= adj_ld_block2;
@@ -298,15 +298,23 @@ status_t brgemm_blocking(brgemm_desc_t *brg) {
                 = (brg->is_f16 && brg->isa_impl == avx512_core_fp16)
                 ? 1
                 : data_type_vnni_granularity(brg->dt_a);
+
         int rd_unroll = one_of(brg->dt_b, data_type::nf4, data_type::u4, data_type::s4, data_type::f4_e2m1) ? 32 : 4;
-        if (brg->with_grouped_wei_decomp) {
+        if (brg->with_grouped_wei_decomp && !brg->with_src_dyn_quant) {
             auto min_group_size = nstl::min(brg->wei_decomp_scales_group_size, brg->wei_decomp_zero_points_group_size);
             min_group_size = nstl::min(min_group_size, brg->src_scales_group_size);
             rd_unroll = nstl::min(rd_unroll, min_group_size / vnni_granularity);
             rd_unroll = nstl::min(rd_unroll, min_group_size / vnni_granularity);
+            brg->rd_block = rd_unroll * vnni_granularity;
+        } else if (brg->with_src_dyn_quant) {
+            brg->rd_block = 32;
+            auto min_group_size = nstl::min(brg->wei_decomp_scales_group_size, brg->wei_decomp_zero_points_group_size);
+            min_group_size = nstl::min(min_group_size, brg->src_scales_group_size);
+            brg->rd_block = nstl::min(brg->rd_block, min_group_size);
+        } else {
+            brg->rd_block = rd_unroll * vnni_granularity;
         }
 
-        brg->rd_block = rd_unroll * vnni_granularity;
         brg->rdb = brg->reduce_dim / brg->rd_block;
         brg->rdb_tail = brg->reduce_dim % brg->rd_block;
 

src/cpu/x64/brgemm/jit_brgemm_kernel.cpp

@@ -203,6 +203,7 @@ struct jit_brgemm_kernel_t : public jit_generator {
     const reg64_t reg_aux_wei_zp = reg_rdb_loop;
     const reg64_t reg_ic = reg_rdb_loop;
     const reg64_t reg_src_scales = reg_rdb_loop;
+    const reg64_t reg_src_grouped_sum = reg_rdb_loop;
     const reg64_t reg_tmp_read_values = reg_rdb_loop;
 
     const reg64_t reg_aux_scales = reg_aux_B;
@@ -280,12 +281,13 @@ struct jit_brgemm_kernel_t : public jit_generator {
     constexpr static int reg_src_scales_offs_ = 336;
     constexpr static int reg_aux_src_scales_offs_ = 344;
     constexpr static int reg_aux2_src_scales_offs_ = 352;
-    // constexpr static int stack_space_needed_ = 360;
+    constexpr static int reg_src_grouped_sum_offs_ = 360;
+    constexpr static int reg_aux_src_grouped_sum_offs_ = 368;
+    constexpr static int reg_aux2_src_grouped_sum_offs_ = 376;
     // these are used for FP8 as temporary push/pop spaces
-    constexpr static int reg_val_tmp_1_ = 368;
-    constexpr static int reg_val_tmp_2_ = 376;
-    constexpr static int stack_space_needed_ = 384;
-    // regsiters for dynamic quant
+    constexpr static int reg_val_tmp_1_ = 384;
+    constexpr static int reg_val_tmp_2_ = 392;
+    constexpr static int stack_space_needed_ = 400;
 
 
     bool is_ldb_loop_ = false;
@@ -323,7 +325,7 @@ struct jit_brgemm_kernel_t : public jit_generator {
         }
 
         if (brg.with_src_dyn_quant) {
-            used_vregs += 2;
+            used_vregs += 1;
         }
 
         if (brg.with_src_dyn_quant && brg.with_wei_decomp_zero_points && brg.wei_decomp_zero_points_stride != 0) {
@@ -971,6 +973,12 @@ void jit_brgemm_kernel_t<Wmm>::copy_post_ops_stack_values_to_aux(
         mov(reg_src_scales, ptr[rsp + reg_src_scales_offs_]);
         mov(ptr[rsp + reg_aux_src_scales_offs_], reg_src_scales);
         mov(ptr[rsp + reg_aux2_src_scales_offs_], reg_src_scales);
+
+        if (brg.with_wei_decomp_zero_points) {
+            mov(reg_src_grouped_sum, ptr[rsp + reg_src_grouped_sum_offs_]);
+            mov(ptr[rsp + reg_aux_src_grouped_sum_offs_], reg_src_grouped_sum);
+            mov(ptr[rsp + reg_aux2_src_grouped_sum_offs_], reg_src_grouped_sum);
+        }
     }
     if (brg.zp_type_b != brgemm_broadcast_t::none) {
         mov(reg_zp_comp_b, ptr[rsp + reg_zp_comp_b_offs_]);
@@ -1048,6 +1056,9 @@ void jit_brgemm_kernel_t<Wmm>::read_params() {
     if (brg.with_src_dyn_quant) {
         mov(reg_src_scales, ptr[param1 + GET_OFF(ptr_src_scales)]);
         mov(ptr[rsp + reg_src_scales_offs_], reg_src_scales);
+
+        mov(reg_src_grouped_sum, ptr[param1 + GET_OFF(ptr_src_grouped_sum)]);
+        mov(ptr[rsp + reg_src_grouped_sum_offs_], reg_src_grouped_sum);
     }
 
     if (brg.zp_type_c != brgemm_broadcast_t::none) {
@@ -2296,21 +2307,10 @@ void jit_brgemm_kernel_t<Wmm>::gemm_microkernel_dyn_quant(int bd_block2,
     };
 
     auto vmm_zero_point = [&](int ld) {
-        int idx = isa_num_vregs(brg.isa_impl) - 3 - ld;
+        int idx = isa_num_vregs(brg.isa_impl) - 2 - ld;
         return Vmm(idx);
     };
 
-    static const int8_t negative_one[64] = {
-        -1, -1, -1, -1, -1, -1, -1, -1,
-        -1, -1, -1, -1, -1, -1, -1, -1,
-        -1, -1, -1, -1, -1, -1, -1, -1,
-        -1, -1, -1, -1, -1, -1, -1, -1,
-        -1, -1, -1, -1, -1, -1, -1, -1,
-        -1, -1, -1, -1, -1, -1, -1, -1,
-        -1, -1, -1, -1, -1, -1, -1, -1,
-        -1, -1, -1, -1, -1, -1, -1, -1
-    };
-
     static const int8_t mask_low_half[64] = {
         0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F,
         0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F,
@@ -2328,33 +2328,18 @@ void jit_brgemm_kernel_t<Wmm>::gemm_microkernel_dyn_quant(int bd_block2,
     if (brg.with_wei_decomp_zero_points) {
         mov(reg_local_wei_zp, ptr[rsp + reg_aux2_wei_zero_points_offs_]);
         if (brg.wei_decomp_zero_points_stride == 0) {
-            auto reg_ptr_8 = Reg8(reg_ptr.getIdx());
-            mov(reg_ptr_8, ptr[reg_local_wei_zp]);
-            uni_vpbroadcastb(vmm_zero_point(0), reg_ptr_8);
+            auto reg_ptr_32 = Reg32(reg_ptr.getIdx());
+            movzx(reg_ptr_32, ptr[reg_local_wei_zp]);
+            uni_vmovq(Xmm(vmm_zero_point(0).getIdx()), reg_ptr);
+            uni_vbroadcastss(vmm_zero_point(0), Xmm(vmm_zero_point(0).getIdx()));
         } else {
-            static const int8_t index_table[64] = {
-                0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x0C, 0x0C, 0x0C, 0x0C,
-                0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x0C, 0x0C, 0x0C, 0x0C,
-                0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x0C, 0x0C, 0x0C, 0x0C,
-                0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x0C, 0x0C, 0x0C, 0x0C
-            };
-
-            auto vmm_indexes = Vmm(isa_num_vregs(brg.isa_impl) - 1);
-            mov(reg_ptr, (size_t)index_table);
-            uni_vmovups(vmm_indexes, ptr[reg_ptr]);
-
             for (int ld = 0; ld < ld_block2; ld++) {
                 uni_vpmovzxbd(vmm_zero_point(ld), ptr[reg_local_wei_zp + ld * brg.ld_block * types::data_type_size(brg.wei_decomp_zero_points_dt)]);
-                vpshufb(vmm_zero_point(ld), vmm_zero_point(ld), vmm_indexes);
             }
         }
     }
 
-    auto vmm_neg_one = Vmm(isa_num_vregs(brg.isa_impl) - 1);
-    mov(reg_ptr, (size_t)negative_one);
-    uni_vmovups(vmm_neg_one, ptr[reg_ptr]);
-
-    auto vmm_mask_low_half = Vmm(isa_num_vregs(brg.isa_impl) - 2);
+    auto vmm_mask_low_half = Vmm(isa_num_vregs(brg.isa_impl) - 1);
     mov(reg_ptr, (size_t)mask_low_half);
     uni_vmovups(vmm_mask_low_half, ptr[reg_ptr]);
 
@@ -2409,22 +2394,28 @@ void jit_brgemm_kernel_t<Wmm>::gemm_microkernel_dyn_quant(int bd_block2,
                 auto vmm = accm(ld_block2, bd, ld);
                 vpdpbusd(vmm, load(ld), bcst(), is_superset(brg.isa_impl, avx512_core) ? EvexEncoding : VexEncoding);
             }
-            if (brg.with_wei_decomp_zero_points) {
-                uni_vpxor(bcst(), bcst(), vmm_neg_one);
-                uni_vpsubb(bcst(), bcst(), vmm_neg_one);
-                for (int ld = 0; ld < ld_block2; ld++) {
-                    auto vmm =  accm(ld_block2, bd, ld);
-                    Vmm vmm_zp = brg.wei_decomp_zero_points_stride == 0 ? vmm_zero_point(0) : vmm_zero_point(ld);
-                    vpdpbusd(vmm, vmm_zp, bcst(), is_superset(brg.isa_impl, avx512_core) ? EvexEncoding : VexEncoding);
-                }
-            }
         }
     }
 
     auto reg_local_src_scales = reg_local_wei_zp;
+    auto reg_local_src_grouped_sum = reg_local_wei_zp;
     auto vmm_src_scales = bcst();
-    mov(reg_local_src_scales, ptr[rsp + reg_aux2_src_scales_offs_ + accums_stack_space]);
+    auto vmm_src_grouped_sum = bcst();
 
+    if (brg.with_wei_decomp_zero_points) {
+        mov(reg_local_src_grouped_sum, ptr[rsp + reg_aux2_src_grouped_sum_offs_ + accums_stack_space]);
+        for (int bd = bd_b; bd < bd_e; bd++) {
+            for (int ld = 0; ld < ld_block2; ld++) {
+                auto vmm_accm = accm(ld_block2, bd, ld);
+                Vmm vmm_zp = brg.wei_decomp_zero_points_stride == 0 ? vmm_zero_point(0) : vmm_zero_point(ld);
+                uni_vbroadcastss(vmm_src_grouped_sum, ptr[reg_local_src_grouped_sum + bd * brg.src_grouped_sum_stride * sizeof(int32_t)]);
+                uni_vpmulld(vmm_src_grouped_sum, vmm_src_grouped_sum, vmm_zp);
+                uni_vpsubd(vmm_accm, vmm_accm, vmm_src_grouped_sum);
+            }
+        }
+    }
+
+    mov(reg_local_src_scales, ptr[rsp + reg_aux2_src_scales_offs_ + accums_stack_space]);
     for (int bd = bd_b; bd < bd_e; bd++) {
         uni_vbroadcastss(vmm_src_scales, ptr[reg_local_src_scales + bd * brg.src_scales_stride * sizeof(float)]);
         for (int ld = 0; ld < ld_block2; ld++) {
@@ -3073,6 +3064,11 @@ void jit_brgemm_kernel_t<Wmm>::ldb_loop(int bd_block2, bool is_bdb_tail,
                         if (brg.with_src_dyn_quant) {
                             ic_group_shift(reg_aux_src_scales_offs_, reg_aux2_src_scales_offs_,
                                            brg.src_scales_group_size, sizeof(float));
+
+                            if (brg.with_wei_decomp_zero_points) {
+                                ic_group_shift(reg_aux_src_grouped_sum_offs_, reg_aux2_src_grouped_sum_offs_,
+                                            brg.src_sum_group_size, sizeof(int32_t));
+                            }
                         }
 
                         mov(reg_local_ic, ptr[rsp + reg_aux_ic_offs_]);
@@ -3306,6 +3302,10 @@ void jit_brgemm_kernel_t<Wmm>::bdb_loop() {
             mov(reg_src_scales, ptr[rsp + reg_src_scales_offs_]);
             add(reg_src_scales, bd_block2 * brg.bd_block * brg.src_scales_stride * sizeof(float));
             mov(ptr[rsp + reg_src_scales_offs_], reg_src_scales);
+
+            mov(reg_src_grouped_sum, ptr[rsp + reg_src_grouped_sum_offs_]);
+            add(reg_src_grouped_sum, bd_block2 * brg.bd_block * brg.src_grouped_sum_stride * sizeof(int32_t));
+            mov(ptr[rsp + reg_src_grouped_sum_offs_], reg_src_grouped_sum);
         }
 
         advance_bd_block2_post_op_regs(bd_block2);