Backport: cpu: x64: improve performance for brgconv w/zp that using large buffer

xuxinzen · web-flow · commit dbb848437d2f · 2025-10-02T18:05:30.000-04:00
diff --git a/src/cpu/x64/jit_brgemm_conv.cpp b/src/cpu/x64/jit_brgemm_conv.cpp
@@ -839,11 +839,15 @@ inline int brgemm_convolution_fwd_t<isa>::get_comp_offset(const int g,
     if (!jcp.src_zero_point && !jcp.s8s8_compensation_required) return 0;
 
     const auto cur_comp_oh = get_comp_oh(oh);
+    const auto cur_owb = ow / jcp.ow_block;
+    const auto cur_comp_ow = jcp.req_cal_comp_pad && jcp.exec_type != exec_vpad
+            ? comp_owb[cur_owb] + ow % jcp.ow_block
+            : ow;
     const auto comp_idx
             = get_comp_ker_idx(kd_b, kd_e, kh_b, kh_e, kw_b, kw_e, cur_comp_oh);
     assert(IMPLICATION(jcp.req_cal_comp_pad, comp_idx >= 0));
     return jcp.req_cal_comp_pad ? g * comp_ocb_sz + ocb * comp_ker_sz
-                    + comp_idx * comp_kw_sz + ow * comp_ow_sz
+                    + comp_idx * comp_kw_sz + cur_comp_ow * comp_ow_sz
                                 : (g * jcp.nb_oc + ocb) * jcp.oc_block;
 }
 
@@ -902,9 +906,8 @@ status_t brgemm_convolution_fwd_t<isa>::init(engine_t *engine) {
     dst_w_sz = static_cast<dim_t>(OW) * jcp.oc_without_padding;
     dst_h_sz = OH * dst_w_sz;
 
-    const auto comp_buffer_os = jcp.exec_type != exec_vpad ? jcp.ow : 1;
     comp_ow_sz = static_cast<dim_t>(jcp.oc_block);
-    comp_kw_sz = comp_buffer_os * comp_ow_sz;
+    comp_kw_sz = comp_ow_sz * jcp.comp_ow_size;
     comp_ker_sz = jcp.ker_ranges_size * comp_kw_sz;
     comp_ocb_sz = jcp.nb_oc * comp_ker_sz;
 
@@ -1161,6 +1164,60 @@ status_t brgemm_convolution_fwd_t<isa>::init(engine_t *engine) {
         }
     }
 
+    if (jcp.req_cal_comp_pad && jcp.exec_type != exec_vpad) {
+        comp_owb.resize(jcp.nb_ow, -1);
+        vector<int> ow_kw_b(jcp.ow, -1);
+        vector<int> ow_kw_e(jcp.ow, -1);
+        vector<int> comp_ow_kw_s(jcp.comp_ow_size, -1);
+        vector<int> comp_ow_kw_f(jcp.comp_ow_size, -1);
+        dim_t comp_ow_l = 0;
+
+        for (int ow = 0; ow < OW;) {
+            assert(comp_ow_l <= jcp.comp_ow_size);
+            const auto iiw = ow * SW - LP;
+            const auto kw_s = div_up(nstl::max(0, -iiw), DW);
+            const auto kw_f = KW
+                    - div_up(nstl::max(0, iiw - IW + (KW - 1) * DW + 1), DW);
+
+            int ow_e = ow;
+            while (ow_e < OW) {
+                const auto iiw_e = ow_e * SW - LP;
+                const auto cur_kw_s = div_up(nstl::max(0, -iiw_e), DW);
+                const auto cur_kw_f = KW
+                        - div_up(nstl::max(0, iiw_e - IW + (KW - 1) * DW + 1),
+                                DW);
+                ow_kw_b[ow_e] = kw_s;
+                ow_kw_e[ow_e] = kw_f;
+                if (cur_kw_s != kw_s || cur_kw_f != kw_f) break;
+                if (ow_e - ow < jcp.ow_block) {
+                    comp_ow_kw_s[comp_ow_l] = kw_s;
+                    comp_ow_kw_f[comp_ow_l] = kw_f;
+                    comp_ow_l++;
+                }
+                ow_e++;
+            }
+            ow = ow_e;
+        }
+
+        for (int owb = 0; owb < jcp.nb_ow; owb++) {
+            const auto ow_b = owb * jcp.ow_block;
+            const auto ow_e = nstl::min(ow_b + jcp.ow_block, OW);
+
+            if (owb == 0) {
+                comp_owb[owb] = 0;
+                continue;
+            }
+
+            for_(int i = 0; i < comp_ow_l; i++)
+            for (int ow = ow_b; ow < ow_e; ow++) {
+                if (ow_kw_b[ow] != comp_ow_kw_s[i + ow - ow_b]
+                        || ow_kw_e[ow] != comp_ow_kw_f[i + ow - ow_b])
+                    break;
+                if (ow == ow_e - 1) comp_owb[owb] = i;
+            }
+        }
+    }
+
     // pre-calculate unique kernel combination
     if (jcp.req_cal_comp_pad) {
         std::set<std::vector<int>> unique_kernels;
@@ -1611,10 +1668,12 @@ void brgemm_convolution_fwd_t<isa>::perform_outwork(
         if (is_postwork) {
             p.apply_comp = has_postcomp;
             p.a_zp_compensation = has_postcomp && jcp.src_zero_point
-                    ? &btc.src_zp_comp_ptr[comp_ker_offs + ow_pw_s * comp_ow_sz]
+                    ? &btc.src_zp_comp_ptr[comp_ker_offs
+                            + (ow_pw_s - ow) * comp_ow_sz]
                     : btc.src_zp_comp_ptr;
             p.s8s8_compensation = has_postcomp && jcp.s8s8_compensation_required
-                    ? &btc.s8s8_comp_ptr[comp_ker_offs + ow_pw_s * comp_ow_sz]
+                    ? &btc.s8s8_comp_ptr[comp_ker_offs
+                            + (ow_pw_s - ow) * comp_ow_sz]
                     : btc.s8s8_comp_ptr;
 
             p.ptr_out = dst_base
@@ -2121,7 +2180,7 @@ void brgemm_convolution_fwd_t<isa>::ker_base(brgemm_thread_ctx_t &btc) const {
         }
 
         const auto post_comp_ker_offs = get_comp_offset(
-                btc.g, btc.ocb, 0, 0, kd_s, kd_f, kh_s, kh_f, 0, KW);
+                btc.g, btc.ocb, 0, ow, kd_s, kd_f, kh_s, kh_f, 0, KW);
         perform_outwork(btc, dst_base, bias_w, ow, g_oc, is_oc_tail, ow_b, ow_e,
                 kd_l, kh_l, do_init, do_postwork, post_comp_ker_offs,
                 do_post_comp);
diff --git a/src/cpu/x64/jit_brgemm_conv.hpp b/src/cpu/x64/jit_brgemm_conv.hpp
@@ -245,7 +245,7 @@ struct brgemm_convolution_fwd_t : public primitive_t {
     std::vector<dim_t> owb_kw_top_vpads;
     std::vector<dim_t> owb_kw_bottom_vpads;
     std::vector<dim_t> kd_bs, kd_es, kh_bs, kh_es, kw_bs, kw_es, oh_kh_b,
-            oh_kh_e, comp_oh, comp_oh_kh_b, comp_oh_kh_e;
+            oh_kh_e, comp_oh, comp_oh_kh_b, comp_oh_kh_e, comp_owb;
 
     int KD, KH, KW, EXT_KD, EXT_KH, EXT_KW, KS, KD_BLOCK, KH_BLOCK, KW_BLOCK,
             KD_BLOCK_PAD, KH_BLOCK_PAD, ID, IH, IW, IDP, IHP, IWP, OD, OH, OW,
diff --git a/src/cpu/x64/jit_brgemm_conv_comp_pad_kernel.cpp b/src/cpu/x64/jit_brgemm_conv_comp_pad_kernel.cpp
@@ -56,7 +56,7 @@ jit_uni_brgemm_conv_comp_pad_kernel_t<Vmm>::
     , out_ker_sz_(static_cast<size_t>(out_ow_sz_)
               * (jcp_.exec_type == exec_trans ? jcp_.prop_kind == backward_data
                                       ? jcp_.iw
-                                      : jcp_.ow
+                                      : jcp_.comp_ow_size
                                               : 1))
     , isa_max_regs(isa_num_vregs(jcp_.isa)) {}
 
@@ -335,27 +335,48 @@ void jit_uni_brgemm_conv_comp_pad_kernel_t<Vmm>::fwd_kw_ow_loop(const int icb,
         const int mb_tail, const int n_block, const bool use_inversion) {
     vector<int> kw_ow_b(jcp_.kw, -1);
     vector<int> kw_ow_e(jcp_.kw, -1);
+    vector<int> comp_ow_kw_s(jcp_.comp_ow_size, -1);
+    vector<int> comp_ow_kw_f(jcp_.comp_ow_size, -1);
+    dim_t comp_ow_l = 0;
     const auto DW = jcp_.dilate_w + 1;
 
-    for (int ow = 0; ow < jcp_.ow; ow++) {
+    for (int ow = 0; ow < jcp_.ow;) {
         const auto iiw = ow * jcp_.stride_w - jcp_.l_pad;
         const auto kw_s = div_up(nstl::max(0, -iiw), DW);
         const auto kw_f = jcp_.kw
                 - div_up(nstl::max(0, iiw - jcp_.iw + (jcp_.kw - 1) * DW + 1),
                         DW);
-        for (int kw = 0; kw < jcp_.kw; kw++) {
-            if (kw >= kw_s && kw < kw_f) {
-                const auto inv_kw = use_inversion ? jcp_.kw - 1 - kw : kw;
-                kw_ow_b[inv_kw] = kw_ow_b[inv_kw] == -1 ? ow : kw_ow_b[inv_kw];
-                kw_ow_e[inv_kw] = ow + 1;
+        int ow_e = ow;
+        while (ow_e < jcp_.ow) {
+            const auto iiw_e = ow_e * jcp_.stride_w - jcp_.l_pad;
+            const auto cur_kw_s = div_up(nstl::max(0, -iiw_e), DW);
+            const auto cur_kw_f = jcp_.kw
+                    - div_up(nstl::max(0,
+                                     iiw_e - jcp_.iw + (jcp_.kw - 1) * DW + 1),
+                            DW);
+            if (cur_kw_s != kw_s || cur_kw_f != kw_f) break;
+            if (ow_e - ow < jcp_.ow_block) {
+                comp_ow_kw_s[comp_ow_l] = kw_s;
+                comp_ow_kw_f[comp_ow_l] = kw_f;
+                comp_ow_l++;
             }
+            ow_e++;
+        }
+        ow = ow_e;
+    }
+
+    for_(int ow = 0; ow < comp_ow_l; ow++)
+    for (int kw = 0; kw < jcp_.kw; kw++) {
+        if (kw >= comp_ow_kw_s[ow] && kw < comp_ow_kw_f[ow]) {
+            const auto inv_kw = use_inversion ? jcp_.kw - 1 - kw : kw;
+            kw_ow_b[inv_kw] = kw_ow_b[inv_kw] == -1 ? ow : kw_ow_b[inv_kw];
+            kw_ow_e[inv_kw] = ow + 1;
         }
     }
 
     for (int kw = 0; kw < jcp_.kw; kw++) {
         const auto ow_b = kw_ow_b[kw];
         const auto ow_e = kw_ow_e[kw];
-
         if (ow_b < ow_e) {
             zero_accumulators(m_block, n_block);
             kdh_loop(icb, icb_tail, ic_step, m_block, mb_tail, n_block);
@@ -366,7 +387,7 @@ void jit_uni_brgemm_conv_comp_pad_kernel_t<Vmm>::fwd_kw_ow_loop(const int icb,
                                                 : inp_kw_sz_);
     }
 
-    copy_ow(m_block, n_block, 0, jcp_.ow);
+    copy_ow(m_block, n_block, 0, comp_ow_l);
 }
 template <typename Vmm>
 void jit_uni_brgemm_conv_comp_pad_kernel_t<Vmm>::kw_loop_base(const int icb,
@@ -590,6 +611,7 @@ void jit_uni_brgemm_conv_relo_comp_pad_kernel_t<Vmm>::kw_loop(
         const int n_block) {
     vector<int> ow_kw_b(jcp_.ow, -1);
     vector<int> ow_kw_e(jcp_.ow, -1);
+    int prev_comp_ow = 0;
     const auto DW = jcp_.dilate_w + 1;
     for (int ow = 0; ow < jcp_.ow; ow++) {
         const auto iiw = ow * jcp_.stride_w - jcp_.l_pad;
@@ -609,12 +631,15 @@ void jit_uni_brgemm_conv_relo_comp_pad_kernel_t<Vmm>::kw_loop(
             if (ow_kw_b[ow_e] != kw_b || ow_kw_e[ow_e] != kw_e) break;
             ow_e++;
         }
+        const auto ow_l = nstl::min(ow_e - ow, jcp_.ow_block);
+
         if (kw_b < kw_e) {
             zero_accumulators(n_block);
             compute(n_block, kw_b, kw_e);
-            store(n_block, ow, ow_e);
+            store(n_block, prev_comp_ow, prev_comp_ow + ow_l);
         }
         ow = ow_e;
+        prev_comp_ow += ow_l;
     }
 }
 
@@ -707,7 +732,7 @@ jit_uni_brgemm_conv_relo_comp_pad_kernel_t<Vmm>::
               * (jcp_.is_relo_whi() ? 1 : jcp_.kw))
     , inp_oc_sz_(static_cast<size_t>(inp_ic_sz_) * jcp_.ic * jcp_.kh * jcp_.kw)
     , out_ow_sz_(static_cast<size_t>(out_dsz_) * jcp_.oc_block)
-    , out_ker_sz_(static_cast<size_t>(out_ow_sz_) * jcp_.ow)
+    , out_ker_sz_(static_cast<size_t>(out_ow_sz_) * jcp_.comp_ow_size)
     , isa_max_regs_(isa_num_vregs(jcp_.isa)) {}
 
 template struct jit_uni_brgemm_conv_comp_pad_kernel_t<Xbyak::Zmm>;
diff --git a/src/cpu/x64/jit_brgemm_conv_utils.cpp b/src/cpu/x64/jit_brgemm_conv_utils.cpp
@@ -205,6 +205,37 @@ struct brg_blocking_t : public jit_brgemm_conv_conf_t {
     loop_t loop[MAXNLOOPS];
 };
 
+dim_t get_comp_ow_size(const jit_brgemm_conv_conf_t &jcp) {
+    if (jcp.exec_type == exec_vpad) return 1;
+
+    const auto IW = jcp.iw;
+    const auto DW = jcp.dilate_w + 1;
+    const auto KW = jcp.kw;
+    const auto SW = jcp.stride_w;
+    const auto LP = jcp.l_pad;
+    dim_t comp_ow_size = 0;
+
+    for (int ow = 0; ow < jcp.ow;) {
+        const auto iiw = ow * SW - LP;
+        const auto kw_s = div_up(nstl::max(0, -iiw), DW);
+        const auto kw_f
+                = KW - div_up(nstl::max(0, iiw - IW + (KW - 1) * DW + 1), DW);
+        int ow_e = ow;
+        while (ow_e < jcp.ow) {
+            const auto iiw_e = ow_e * SW - LP;
+            const auto cur_kw_s = div_up(nstl::max(0, -iiw_e), DW);
+            const auto cur_kw_f = KW
+                    - div_up(nstl::max(0, iiw_e - IW + (KW - 1) * DW + 1), DW);
+            if (cur_kw_s != kw_s || cur_kw_f != kw_f) break;
+            if (ow_e - ow < jcp.ow_block) comp_ow_size++;
+            ow_e++;
+        }
+        ow = ow_e;
+    }
+
+    return comp_ow_size;
+}
+
 bool is_any_eligible(const jit_brgemm_conv_conf_t &jcp) {
     return (jcp.prop_kind == prop_kind::forward_inference || jcp.wei_plain
             || one_of(jcp.wei_dt, data_type::s8, data_type::f16,
@@ -2240,9 +2271,11 @@ status_t init_conf(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa,
                 && IMPLICATION(jcp.wei_dt == f16, isa != avx10_1_512)
                 && jcp.ic * rd_ksize > rd_padded_block;
 
+        // Disable os blocking to avoid using large buffer
+        // The value is empirical
         jcp.is_os_blocking = jcp.f_pad < jcp.kd && jcp.back_pad < jcp.kd
                 && jcp.t_pad < jcp.kh && jcp.b_pad < jcp.kh
-                && jcp.r_pad < jcp.kw && jcp.l_pad < jcp.kw;
+                && jcp.r_pad < jcp.kw && jcp.l_pad < jcp.kw && jcp.iwp < 10800;
 
         if (is_amx(isa)
                 && IMPLICATION(!jcp.is_relo(),
@@ -2397,22 +2430,15 @@ status_t init_conf(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa,
             + utils::div_up(abs(jcp.back_pad), jcp.dilate_d + 1);
     const auto kh_cnt = 1 + utils::div_up(abs(jcp.t_pad), jcp.dilate_h + 1)
             + utils::div_up(abs(jcp.b_pad), jcp.dilate_h + 1);
-    jcp.ker_ranges_size = jcp.exec_type == exec_trans
-            ? kd_cnt * nstl::min(jcp.oh, jcp.oh_block + kh_cnt)
-            : kd_cnt * kh_cnt;
-    const auto comp_buffer_ow = jcp.exec_type != exec_vpad ? jcp.ow : 1;
+    jcp.ker_ranges_size = jcp.exec_type == exec_trans ? kd_cnt
+                    * nstl::min(
+                            jcp.oh, rnd_up(jcp.oh_block + kh_cnt, jcp.oh_block))
+                                                      : kd_cnt * kh_cnt;
+    jcp.comp_ow_size = get_comp_ow_size(jcp);
     jcp.comp_a_buffer_size = jcp.ngroups * jcp.nb_oc * jcp.ker_ranges_size
-            * comp_buffer_ow * jcp.oc_block;
+            * jcp.comp_ow_size * jcp.oc_block;
     jcp.s8s8_comp_buffer_size = jcp.comp_a_buffer_size;
 
-    // Dispatch the shapes to VNNI for better performance
-    // TODO: optimize the perf for zero point with large buffer on AMX
-    if (is_amx(isa) && jcp.src_zero_point && jcp.exec_type == exec_trans
-            && (jcp.l_pad > 0 || jcp.r_pad > 0) && jcp.oc * jcp.ow > 8192)
-        VDISPATCH_CONV_IC(!allow_perf_heuristics(jcp),
-                VERBOSE_IMPL_HEURISTIC_FAIL,
-                "no optimization for zero point on amx")
-
     // For padding shapes, we calculate the comp along with the computation
     // inside brgemm kernel when output size is small to get optimal perf
     // For shapes with large ow we calculate the comp inside brgemm kernel too
@@ -2653,6 +2679,7 @@ status_t init_1x1_conf(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa,
         weights_md.extra.asymm_compensation_mask = with_groups ? 0x3 : 0x1;
     }
     jcp.req_cal_comp_pad = false;
+    jcp.comp_ow_size = 1;
     jcp.s8s8_comp_buffer_size = jcp.ngroups * jcp.nb_oc * jcp.oc_block;
     jcp.comp_a_buffer_size = jcp.ngroups * jcp.nb_oc * jcp.oc_block;
 
diff --git a/src/cpu/x64/jit_primitive_conf.hpp b/src/cpu/x64/jit_primitive_conf.hpp
@@ -780,6 +780,7 @@ struct jit_brgemm_conv_conf_t {
 
     bool use_buffer;
     dim_t buffer_size;
+    dim_t comp_ow_size;
     dim_t ker_ranges_size;
     dim_t comp_a_buffer_size;
     dim_t s8s8_comp_buffer_size;
