[FORK][FEATURE] DQ IP: reduce aux vecs counts required for microkernel

Author: dmitrygo

src/cpu/x64/brgemm/brgemm_utils.cpp

@@ -230,9 +230,8 @@ int calculate_max_bcast_block(brgemm_desc_t *brg, const int adj_ld_block2) {
     if (one_of(brg->dt_b, data_type::nf4) && brg->isa_impl == avx2) max_bcast_block -= 5;
     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_wei_decomp_zero_points && brg->wei_decomp_zero_points_stride == 0 && !brg->with_src_dyn_quant) max_bcast_block -= 1;
     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;
 

src/cpu/x64/brgemm/jit_brgemm_kernel.cpp

@@ -320,17 +320,13 @@ struct jit_brgemm_kernel_t : public jit_generator {
             used_vregs += 1;
         }
 
-        if (brg.with_wei_decomp_zero_points && brg.wei_decomp_zero_points_stride == 0) {
+        if (brg.with_wei_decomp_zero_points && brg.wei_decomp_zero_points_stride == 0 && !brg.with_src_dyn_quant) {
             used_vregs += 1;
         }
 
         if (brg.with_src_dyn_quant) {
             used_vregs += 1;
         }
-
-        if (brg.with_src_dyn_quant && brg.with_wei_decomp_zero_points && brg.wei_decomp_zero_points_stride != 0) {
-            used_vregs += brg.ld_block2;
-        }
         return isa_num_vregs(brg.isa_impl) - used_vregs;
     }
 
@@ -2306,11 +2302,6 @@ void jit_brgemm_kernel_t<Wmm>::gemm_microkernel_dyn_quant(int bd_block2,
         if (brg.req_s8s8_compensation) uni_vpaddb(v1, v1, vmm_inp_shift());
     };
 
-    auto vmm_zero_point = [&](int ld) {
-        int idx = isa_num_vregs(brg.isa_impl) - 2 - ld;
-        return Vmm(idx);
-    };
-
     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,
@@ -2321,30 +2312,11 @@ void jit_brgemm_kernel_t<Wmm>::gemm_microkernel_dyn_quant(int bd_block2,
     mov(ptr[rsp + reg_bdb_loop_offs_], reg_bdb_loop);
     mov(ptr[rsp + reg_ldb_loop_offs_], reg_ldb_loop);
 
-    auto reg_local_wei_scales = reg_bdb_loop;
-    auto reg_local_wei_zp = reg_ldb_loop;
-    auto reg_ptr = reg_local_wei_scales;
-
-    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_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 {
-            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)]);
-            }
-        }
-    }
-
+    auto reg_ptr = reg_bdb_loop;
     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]);
 
-    mov(reg_local_wei_scales, ptr[rsp + reg_aux2_wei_scales_offs_]);
-
     const int vec_size = vreg_traits<Vmm>::vlen;
     auto accums_stack_space = bd_e * ld_block2 * vec_size;
     sub(rsp, accums_stack_space);
@@ -2397,42 +2369,68 @@ void jit_brgemm_kernel_t<Wmm>::gemm_microkernel_dyn_quant(int bd_block2,
         }
     }
 
-    auto reg_local_src_scales = reg_local_wei_zp;
-    auto reg_local_src_grouped_sum = reg_local_wei_zp;
-    auto vmm_src_scales = bcst();
+    auto vmm_zero_point = [&](int ld) {
+        return load(ld);
+    };
+
+    auto reg_local_wei_zp = reg_ldb_loop;
+    auto reg_local_src_grouped_sum = reg_bdb_loop;
+    auto vmm_tmp = Vmm(isa_num_vregs(brg.isa_impl) - 1);
     auto vmm_src_grouped_sum = bcst();
 
     if (brg.with_wei_decomp_zero_points) {
+        mov(reg_local_wei_zp, ptr[rsp + reg_aux2_wei_zero_points_offs_ + accums_stack_space]);
+        if (brg.wei_decomp_zero_points_stride == 0) {
+            Vmm vmm_zp = vmm_zero_point(0);
+            auto reg_ptr_32 = Reg32(reg_ptr.getIdx());
+            movzx(reg_ptr_32, ptr[reg_local_wei_zp]);
+            uni_vmovq(Xmm(vmm_zp.getIdx()), reg_ptr);
+            uni_vbroadcastss(vmm_zp, Xmm(vmm_zp.getIdx()));
+        }
+
         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++) {
+            uni_vbroadcastss(vmm_src_grouped_sum, ptr[reg_local_src_grouped_sum + bd * brg.src_grouped_sum_stride * sizeof(int32_t)]);
             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);
+                if (bd == bd_b && brg.wei_decomp_zero_points_stride != 0) {
+                    uni_vpmovzxbd(vmm_zp, ptr[reg_local_wei_zp + ld * brg.ld_block * types::data_type_size(brg.wei_decomp_zero_points_dt)]);
+                }
+
+                auto vmm_accm = accm(ld_block2, bd, ld);
+                uni_vpmulld(vmm_tmp, vmm_src_grouped_sum, vmm_zp);
+                uni_vpsubd(vmm_accm, vmm_accm, vmm_tmp);
             }
         }
     }
 
+    auto wei_scale = [&](int ld) {
+        return load(ld);
+    };
+
+    auto reg_local_src_scales = reg_ldb_loop;
+    auto reg_local_wei_scales = reg_bdb_loop;
+    auto vmm_src_scales = bcst();
+
+    mov(reg_local_wei_scales, ptr[rsp + reg_aux2_wei_scales_offs_ + accums_stack_space]);
     mov(reg_local_src_scales, ptr[rsp + reg_aux2_src_scales_offs_ + accums_stack_space]);
+    if (brg.wei_decomp_scales_stride == 0) {
+        uni_vbroadcastss(wei_scale(0), ptr[reg_local_wei_scales]);
+    }
+
     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++) {
-            if (brg.wei_decomp_scales_stride == 0) {
-                uni_vbroadcastss(load(ld), ptr[reg_local_wei_scales]);
-            } else {
-                uni_vmovups(load(ld), ptr[reg_local_wei_scales + ld * brg.ld_block * sizeof(float)]);
+            auto vmm_wei_scale = brg.wei_decomp_scales_stride == 0 ? wei_scale(0) : wei_scale(ld);
+            if (bd == bd_b && brg.wei_decomp_scales_stride != 0) {
+                uni_vmovups(vmm_wei_scale, ptr[reg_local_wei_scales + ld * brg.ld_block * sizeof(float)]);
             }
-        }
-        for (int ld = 0; ld < ld_block2; ld++) {
-            auto vmm_accm = accm(ld_block2, bd, ld);
 
+            auto vmm_accm = accm(ld_block2, bd, ld);
             uni_vcvtdq2ps(vmm_accm, vmm_accm);
-            uni_vmulps(vmm_accm, vmm_accm, vmm_src_scales);
-            uni_vmulps(load(ld), vmm_accm, load(ld));
+            uni_vmulps(vmm_tmp, vmm_accm, vmm_src_scales);
             uni_vmovups(vmm_accm, ptr[rsp + (bd * ld_block2 + ld) * vec_size]);
-            uni_vaddps(vmm_accm, vmm_accm, load(ld));
+            uni_vfmadd231ps(vmm_accm, vmm_tmp, vmm_wei_scale);
         }
     }