cpu: x64: brg_matmul: postpone access zp buffers before parallel task

dzarukin · dzarukin · commit 327fbbf975e4 · 2025-11-19T14:21:46.000-08:00
diff --git a/src/cpu/x64/matmul/brgemm_matmul.cpp b/src/cpu/x64/matmul/brgemm_matmul.cpp
@@ -749,7 +749,6 @@ void brgemm_matmul_t<isa>::compute_kernel(
             = brgmm_ctx.get_zp_a_compensation_ptr(ithr, b_idx, n_blk_idx);
     const auto zp_comp_b
             = brgmm_ctx.get_zp_b_compensation_result_ptr(ithr, m_blk_idx);
-    const auto zp_c_val_ptr = brgmm_ctx.get_zp_c_val_ptr();
     const auto &post_ops_binary_rhs_arg_vec
             = brgmm_ctx.get_post_ops_binary_rhs_arg_vec();
     const bool post_ops_applicable = bgmmc.post_ops_applicable
@@ -794,8 +793,8 @@ void brgemm_matmul_t<isa>::compute_kernel(
                     first_mb_matrix_addr_off,
                     static_cast<const void *>(zp_comp_a),
                     static_cast<const void *>(zp_comp_b),
-                    static_cast<const void *>(zp_c_val_ptr), false, 1, false,
-                    false, brgmm_ctx.get_src_scales_ptr(),
+                    brgmm_ctx.get_zp_c_ptr(), false, 1, false, false,
+                    brgmm_ctx.get_src_scales_ptr(),
                     brgmm_ctx.get_wei_scales_ptr(n),
                     brgmm_ctx.get_dst_scales_inv_ptr(ithr)};
             brgemm_kernel_execute_postops(brg_kernel, gemm_batch, addr_batch,
@@ -850,8 +849,8 @@ void brgemm_matmul_t<isa>::compute_kernel(
                     first_mb_matrix_addr_off,
                     static_cast<const void *>(zp_comp_a),
                     static_cast<const void *>(zp_comp_b),
-                    static_cast<const void *>(zp_c_val_ptr), false, 1, false,
-                    false, brgmm_ctx.get_src_scales_ptr(),
+                    brgmm_ctx.get_zp_c_ptr(), false, 1, false, false,
+                    brgmm_ctx.get_src_scales_ptr(),
                     brgmm_ctx.get_wei_scales_ptr(n),
                     brgmm_ctx.get_dst_scales_inv_ptr(ithr)};
 
@@ -1135,7 +1134,6 @@ void brgemm_matmul_t<isa>::maybe_reduce_partial_results_and_apply_postops(
                         const auto zp_comp_b
                                 = brgmm_ctx.get_zp_b_compensation_result_ptr(
                                         ithr, mb);
-                        const auto zp_c_val_ptr = brgmm_ctx.get_zp_c_val_ptr();
                         const auto &post_ops_binary_rhs_arg_vec
                                 = brgmm_ctx.get_post_ops_binary_rhs_arg_vec();
 
@@ -1158,9 +1156,8 @@ void brgemm_matmul_t<isa>::maybe_reduce_partial_results_and_apply_postops(
                                 dst_anchor_point, first_mb_matrix_addr_off,
                                 static_cast<const void *>(zp_comp_a),
                                 static_cast<const void *>(zp_comp_b),
-                                static_cast<const void *>(zp_c_val_ptr),
-                                skip_accumulation, 1, false, false,
-                                brgmm_ctx.get_src_scales_ptr(),
+                                brgmm_ctx.get_zp_c_ptr(), skip_accumulation, 1,
+                                false, false, brgmm_ctx.get_src_scales_ptr(),
                                 brgmm_ctx.get_wei_scales_ptr(n),
                                 brgmm_ctx.get_dst_scales_inv_ptr(ithr)};
 
@@ -1199,9 +1196,18 @@ void brgemm_matmul_t<isa>::copy_a_chunk_in_buffer(
     ctx.zp_a_compensation_result_ptr
             = (void *)brgmm_ctx.get_zp_b_compensation_result_ptr(
                     ithr, m_blk_idx);
-    ctx.zp_ab_comp_ptr = (void *)brgmm_ctx.get_zp_ab_mixed_comp_ptr();
     ctx.dynamic_src_ld = brgmm_ctx.get_src_stride();
-    ctx.zp_b_neg_val_ptr = brgmm_ctx.get_wei_zp_neg_ptr();
+
+    // Note: instead of passing an address to a stack variable, a kernel may be
+    // changed to take just zp_b value and perform negation itself, but updating
+    // kernels is not straightforward for all platforms.
+    int32_t neg_zp_b
+            = !bgmmc.with_wei_decompression ? brgmm_ctx.get_neg_zp_b() : 0;
+    int32_t neg_zp_ab_comp = !bgmmc.with_wei_decompression
+            ? bgmmc.K * brgmm_ctx.get_neg_zp_a()
+            : 0;
+    ctx.zp_b_neg_val_ptr = &neg_zp_b;
+    ctx.zp_ab_comp_ptr = &neg_zp_ab_comp;
 
     for (int gb = 0; gb < gemm_batch_iters; gb++) {
         const int k = k_start + gb * bgmmc.K_blk;
@@ -1260,7 +1266,13 @@ void brgemm_matmul_t<isa>::copy_b_chunk_in_buffer(
 
     ctx.zp_a_compensation_ptr = (void *)brgmm_ctx.get_zp_a_compensation_ptr(
             ithr, b_idx, n_blk_idx);
-    ctx.zp_a_neg_value_ptr = (void *)brgmm_ctx.get_zp_a_neg_val_ptr();
+
+    // Note: instead of passing an address to a stack variable, a kernel may be
+    // changed to take just zp_a value and perform negation itself, but updating
+    // kernels is not straightforward for all platforms.
+    int32_t neg_zp_a = brgmm_ctx.get_neg_zp_a();
+    ctx.zp_a_neg_value_ptr = &neg_zp_a;
+
     ctx.compensation_ptr
             = (void *)brgmm_ctx.get_s8s8_comp_ptr(ithr, b_idx, n_blk_idx);
 
@@ -1379,31 +1391,13 @@ struct brgemm_matmul_t<isa>::brg_matmul_exec_ctx_t {
         bias_ptr_ = CTX_IN_MEM(const char *, DNNL_ARG_BIAS);
 
         // setup scales / zp pointers
-        const void *src_zero_points = CTX_IN_MEM(
+        src_zp_ptr_ = CTX_IN_MEM(
                 const void *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC);
         wei_zp_ptr_ = CTX_IN_MEM(
                 const void *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_WEIGHTS);
-        const void *dst_zero_points = CTX_IN_MEM(
+        dst_zp_ptr_ = CTX_IN_MEM(
                 const void *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_DST);
 
-        zero_point_a_negative_val_ = src_zero_points
-                ? -cpu::io::load_int_value(
-                        pd->attr()->zero_points_.get_data_type(DNNL_ARG_SRC),
-                        src_zero_points, 0)
-                : 0;
-        zero_point_c_val_ = dst_zero_points
-                ? cpu::io::load_int_value(
-                        pd->attr()->zero_points_.get_data_type(DNNL_ARG_DST),
-                        dst_zero_points, 0)
-                : 0;
-
-        wei_zp_neg_val_ = (-1)
-                * (wei_zp_ptr_ ? cpu::io::load_int_value(
-                           pd->attr()->zero_points_.get_data_type(
-                                   DNNL_ARG_WEIGHTS),
-                           wei_zp_ptr_, 0)
-                               : 0);
-
         const auto &scratchpad = ctx.get_scratchpad_grantor();
 
         const auto &bgmmc = pd->get_brgemm_matmul_conf();
@@ -1470,9 +1464,6 @@ struct brgemm_matmul_t<isa>::brg_matmul_exec_ctx_t {
                         key_brgemm_primitive_zp_comp_b)
                 : nullptr;
 
-        zero_point_mixed_ab_compensation_component_
-                = bgmmc.K * zero_point_a_negative_val_;
-
         post_ops_binary_rhs_arg_vec_ = binary_injector::prepare_binary_args(
                 pd->attr()->post_ops_, ctx);
         base_brg_ker_idx_
@@ -2100,11 +2091,19 @@ struct brgemm_matmul_t<isa>::brg_matmul_exec_ctx_t {
                 + ithr * sizeof(float);
     }
 
-    const int32_t *get_zp_a_neg_val_ptr() const {
-        return &zero_point_a_negative_val_;
+    int32_t get_neg_zp_a() const {
+        if (!bgmmc_.has_zero_point_a) return 0;
+        return -cpu::io::load_int_value(bgmmc_.src_zp_dt, src_zp_ptr_, 0);
     }
 
-    const void *get_wei_zp_neg_ptr() const { return &wei_zp_neg_val_; }
+    // Used to compute compensation. Can't initialize the value at construction
+    // time as memory buffers must be accessed inside a parallel task
+    // (asynchronous runtime requirement).
+    int32_t get_neg_zp_b() const {
+        if (!bgmmc_.has_zero_point_b) return 0;
+        assert(bgmmc_.is_wei_zp_common);
+        return -cpu::io::load_int_value(bgmmc_.wei_zp_dt, wei_zp_ptr_, 0);
+    }
 
     const void *get_wei_zp_ptr(int n, int k = 0) const {
         if (!bgmmc_.has_zero_point_b) return nullptr;
@@ -2127,11 +2126,7 @@ struct brgemm_matmul_t<isa>::brg_matmul_exec_ctx_t {
         return (char *)wei_zp_ptr_ + offset;
     }
 
-    const int32_t *get_zp_ab_mixed_comp_ptr() const {
-        return &zero_point_mixed_ab_compensation_component_;
-    }
-
-    const int32_t *get_zp_c_val_ptr() const { return &zero_point_c_val_; }
+    const void *get_zp_c_ptr() const { return dst_zp_ptr_; }
 
     int32_t *get_zp_a_compensation_ptr(
             int ithr, int b_idx, int n_blk_idx) const {
@@ -2152,7 +2147,7 @@ struct brgemm_matmul_t<isa>::brg_matmul_exec_ctx_t {
                     + n_blk_idx * bgmmc_.wei_n_blk;
             PRAGMA_OMP_SIMD()
             for (int b = 0; b < bgmmc_.wei_n_blk; b++)
-                zp_comp[b] = -zero_point_a_negative_val_
+                zp_comp[b] = -get_neg_zp_a()
                         * reorder_zp_a_comp_ptr_[base_offset + b];
         }
         return zp_comp;
@@ -2478,11 +2473,9 @@ struct brgemm_matmul_t<isa>::brg_matmul_exec_ctx_t {
     int32_t *zero_point_b_compensations_ptr_;
     int32_t *reorder_zp_a_comp_ptr_;
 
-    int32_t zero_point_a_negative_val_;
-    int32_t zero_point_mixed_ab_compensation_component_;
-    int32_t zero_point_c_val_;
-    int32_t wei_zp_neg_val_;
+    const void *src_zp_ptr_;
     const void *wei_zp_ptr_;
+    const void *dst_zp_ptr_;
     std::vector<const void *> post_ops_binary_rhs_arg_vec_;
 
     int base_brg_ker_idx_;
diff --git a/src/cpu/x64/matmul/brgemm_matmul_copy_utils.hpp b/src/cpu/x64/matmul/brgemm_matmul_copy_utils.hpp
@@ -27,20 +27,20 @@ namespace matmul {
 
 struct jit_brgemm_matmul_copy_b_t {
     struct ctx_t {
-        const void *src;
-        const void *tr_src;
-        const void *compensation_ptr;
-        const void *zp_a_compensation_ptr;
-        const void *zp_a_neg_value_ptr;
-        const void *zp_b_value_ptr;
-        const void *src_scales_ptr;
-        const void *wei_scales_ptr;
-
-        dim_t current_K_start;
-        dim_t current_K_iters;
-        dim_t current_K_pad {0};
-        dim_t current_N_blk;
-        dim_t dynamic_src_stride;
+        const void *src = nullptr;
+        const void *tr_src = nullptr;
+        const void *compensation_ptr = nullptr;
+        const void *zp_a_compensation_ptr = nullptr;
+        const void *zp_a_neg_value_ptr = nullptr;
+        const void *zp_b_value_ptr = nullptr;
+        const void *src_scales_ptr = nullptr;
+        const void *wei_scales_ptr = nullptr;
+
+        dim_t current_K_start = 0;
+        dim_t current_K_iters = 0;
+        dim_t current_K_pad = 0;
+        dim_t current_N_blk = 0;
+        dim_t dynamic_src_stride = 0;
     };
 
     virtual void operator()(ctx_t *ctx) = 0;
@@ -55,17 +55,17 @@ struct jit_brgemm_matmul_copy_b_t {
 
 struct jit_brgemm_matmul_copy_a_t {
     struct ctx_t {
-        const void *src;
-        const void *tr_src;
-        const void *zp_b_compensation_buffer_ptr;
-        const void *zp_a_compensation_result_ptr;
-        const void *zp_b_neg_val_ptr;
-        const void *zp_ab_comp_ptr;
-
-        dim_t current_K_start;
-        dim_t current_K_blk;
-        dim_t current_M_blk;
-        dim_t dynamic_src_ld;
+        const void *src = nullptr;
+        const void *tr_src = nullptr;
+        const void *zp_b_compensation_buffer_ptr = nullptr;
+        const void *zp_a_compensation_result_ptr = nullptr;
+        const void *zp_b_neg_val_ptr = nullptr;
+        const void *zp_ab_comp_ptr = nullptr;
+
+        dim_t current_K_start = 0;
+        dim_t current_K_blk = 0;
+        dim_t current_M_blk = 0;
+        dim_t dynamic_src_ld = 0;
     };
 
     virtual void operator()(ctx_t *ctx) = 0;
diff --git a/src/cpu/x64/matmul/brgemm_matmul_utils.cpp b/src/cpu/x64/matmul/brgemm_matmul_utils.cpp
@@ -1335,7 +1335,6 @@ status_t init_brgemm_matmul_conf(cpu_isa_t isa, brgemm_matmul_conf_t &bgmmc,
     bgmmc.dst_dt = dst_d.data_type();
     bgmmc.wei_dt = weights_d.data_type();
     bgmmc.orig_wei_dt = weights_d.data_type();
-    bgmmc.wei_zp_dt = attr.zero_points_.get(DNNL_ARG_WEIGHTS).get_data_type();
 
     bgmmc.with_reduce = mmd.reduce_desc.format_kind != format_kind::undef;
     bgmmc.reduce_dt
@@ -1466,6 +1465,10 @@ status_t init_brgemm_matmul_conf(cpu_isa_t isa, brgemm_matmul_conf_t &bgmmc,
     VCONDCHECK_BG(!(bgmmc.with_dst_scales && dst_scales.get_mask() > 0),
             VERBOSE_UNSUPPORTED_SCALES_CFG);
 
+    const auto &src_zp = attr.zero_points_.get(DNNL_ARG_SRC);
+    const auto has_src_zp = !src_zp.has_default_values();
+    if (has_src_zp) { bgmmc.src_zp_dt = src_zp.get_data_type(); }
+
     const auto &wei_zp = attr.zero_points_.get(DNNL_ARG_WEIGHTS);
     const auto has_wei_zp = !wei_zp.has_default_values();
 
diff --git a/src/cpu/x64/matmul/brgemm_matmul_utils.hpp b/src/cpu/x64/matmul/brgemm_matmul_utils.hpp
@@ -115,9 +115,6 @@ struct brgemm_matmul_conf_t {
     bool packed_sparse_weights;
     bool with_wei_decompression;
     int postops_inst_count;
-    brgemm_broadcast_t src_zp_type;
-    brgemm_broadcast_t wei_zp_type;
-    brgemm_broadcast_t dst_zp_type;
 
     bool use_buffer_a;
     bool use_buffer_a_tail_only;
@@ -189,7 +186,6 @@ struct brgemm_matmul_conf_t {
     dim_t s8s8_comp_ithr_str;
     dim_t s8s8_comp_b_str;
     dim_t s8s8_comp_n_str;
-    bool has_zero_point_a, has_zero_point_b, has_zero_point_c;
     bool post_ops_applicable;
     bool transposed_A;
     bool transposed_B;
@@ -202,14 +198,6 @@ struct brgemm_matmul_conf_t {
     // were changed.
     bool adjust_a_strides = false;
 
-    dim_t zp_a_comp_shift_n;
-    dim_t zp_a_comp_elems_per_thr;
-
-    dim_t zp_b_comp_result_shift_m;
-    dim_t zp_b_comp_buffer_start;
-    dim_t zp_b_comp_buffer_shift_m;
-    dim_t zp_b_comp_elems_per_thr;
-
     int wsp_tile_per_thr_bytes;
     int brgemm_batch_element_per_thr_sz;
     bool is_amx;
@@ -244,12 +232,29 @@ struct brgemm_matmul_conf_t {
     data_type_t wei_scales_dt = data_type::undef;
 
     // Zero points
+    bool has_zero_point_a;
+    bool has_zero_point_b;
+    bool has_zero_point_c;
+    brgemm_broadcast_t src_zp_type;
+    brgemm_broadcast_t wei_zp_type;
+    brgemm_broadcast_t dst_zp_type;
+
+    data_type_t src_zp_dt = data_type::undef;
+
     dim_t wei_zp_k_gsize = 0;
     bool is_wei_zp_per_k = false;
     bool is_wei_zp_per_n = false;
     bool is_wei_zp_common = false;
     data_type_t wei_zp_dt = data_type::undef;
 
+    dim_t zp_a_comp_shift_n;
+    dim_t zp_a_comp_elems_per_thr;
+
+    dim_t zp_b_comp_result_shift_m;
+    dim_t zp_b_comp_buffer_start;
+    dim_t zp_b_comp_buffer_shift_m;
+    dim_t zp_b_comp_elems_per_thr;
+
     bool is_gemv = false;
     // Currently, it's only used to enable the N=1 code path for M=1, when B
     // is transposed.
