[backport] x64: matmul: Fix nthr issues in matmul (#3248)

inteldimitrius · web-flow · commit ea1b4a169d3f · 2025-05-13T13:24:23.000+05:00
diff --git a/src/cpu/x64/matmul/amx_blocking_heuristics.cpp b/src/cpu/x64/matmul/amx_blocking_heuristics.cpp
@@ -757,7 +757,7 @@ bool matmul_amx_blocking_params_macro_t::set_blocking_parameters() {
     current_lda_ = get_actual_lda();
 
     // Need a temp C buffer if a BRGEMM creates partial results
-    need_buf_c_ = (nthr_k_ != 1) || (k_blk_ != K);
+    need_buf_c_ = (nthr_k_ > 1);
 
     efficiency_score_ = calculate_blocking_scores();
 
@@ -833,6 +833,7 @@ void matmul_amx_blocking_params_micro_t::find_best_blocking(
         const int min_N_blk = !bgmmc.is_runtime_N && low_parallelism
                         && is_amx_xf16 && !bm_conf_utils.check_n_blk_fixed()
                         && bgmmc.N_blk > 32 && !runtime_dims
+                        && !bgmmc.transposed_B // Transposed copy B kernel doesn't support adjusting N_blk
                 ? 32
                 : bgmmc.N_blk;
         const int desired_M_chunk = bgmmc.is_runtime_M
diff --git a/src/cpu/x64/matmul/brgemm_matmul_utils.cpp b/src/cpu/x64/matmul/brgemm_matmul_utils.cpp
@@ -657,13 +657,12 @@ struct matmul_avx512_blocking_params_t {
         return lda;
     }
 
-    inline bool is_buffer_c_required(
-            dim_t acc_dt, dim_t dst_dt, bool with_sum) const {
-        const size_t k_chunk_elems = k_blk * batch_size;
+    inline bool is_buffer_c_required(const brgemm_matmul_conf_t &bgmmc) const {
+        const size_t k_chunk_elems = bgmmc.K_chunk_elems;
         if (nthr_k > 1 && static_cast<size_t>(mp.K) > k_chunk_elems)
             return true;
 
-        return ((acc_dt != dst_dt || with_sum)
+        return ((bgmmc.acc_dt != bgmmc.dst_dt || bgmmc.with_sum)
                 && (static_cast<size_t>(mp.K) > k_chunk_elems
                         || mp.K % k_blk > 0));
     }
@@ -679,8 +678,7 @@ struct matmul_avx512_blocking_params_t {
 
         bgmmc.nthr_k = nthr_k;
 
-        bgmmc.use_buffer_c = is_buffer_c_required(
-                bgmmc.acc_dt, bgmmc.dst_dt, bgmmc.with_sum);
+        bgmmc.use_buffer_c = is_buffer_c_required(bgmmc);
         bgmmc.LDA = bgmmc.adjust_a_strides || bgmmc.use_buffer_a
                         || bgmmc.treat_A_as_plain
                 ? get_actual_lda(bgmmc.use_buffer_a, bgmmc.tr_a_dt_sz)
@@ -1740,21 +1738,12 @@ void init_aux_values(brgemm_matmul_conf_t &bgmmc,
 
     bgmmc.buffer_a_per_thread_sz = bgmmc.buffer_a_m_stride * bgmmc.M_chunk_size;
 
-    bgmmc.buffer_b_gb_stride = bgmmc.tr_b_dt_sz * rnd_up(bgmmc.LDB, bgmmc.N_blk)
-            * bgmmc.wei_k_blk;
+    bgmmc.buffer_b_gb_stride
+            = bgmmc.tr_b_dt_sz * bgmmc.LDB * bgmmc.K_blk * bgmmc.wei_k_blk;
     bgmmc.buffer_b_k_brg_stride
-            = bgmmc.buffer_b_gb_stride * bgmmc.K_chunk_size * bgmmc.K_blk;
-
-    // TODO: (Refactoring) the only usage of `buffer_b_chunk_sz`
-    // and `buffer_b_n_blk_stride` is here. Remove it if not needed.
-    bgmmc.buffer_b_n_blk_stride = bgmmc.tr_b_dt_sz
-            * ((bgmmc.N_blk / bgmmc.LDB) * bgmmc.LDB2
-                    + (bgmmc.N_blk % bgmmc.LDB)
-                            * data_type_vnni_granularity(bgmmc.wei_dt));
-
-    bgmmc.buffer_b_chunk_sz
-            = bgmmc.buffer_b_k_brg_stride * bgmmc.brgemm_batch_size;
-    bgmmc.buffer_b_per_thread_sz = bgmmc.buffer_b_chunk_sz;
+            = bgmmc.buffer_b_gb_stride * bgmmc.brgemm_batch_size;
+    bgmmc.buffer_b_per_thread_sz
+            = bgmmc.buffer_b_k_brg_stride * bgmmc.K_chunk_size;
 
     bgmmc.buffer_reduce_per_thread_sz = 0;
     if (bgmmc.reduce_kind == matmul_reduce_kind::src) {
diff --git a/src/cpu/x64/matmul/brgemm_matmul_utils.hpp b/src/cpu/x64/matmul/brgemm_matmul_utils.hpp
@@ -174,9 +174,6 @@ struct brgemm_matmul_conf_t {
 
     dim_t buffer_b_gb_stride;
     dim_t buffer_b_k_brg_stride;
-    dim_t buffer_b_n_blk_stride;
-
-    dim_t buffer_b_chunk_sz;
     dim_t buffer_b_per_thread_sz;
 
     dim_t buffer_reduce_per_thread_sz;
