cpu: x64: matmul: fix brgemm batch size

Author: yair-obodovsky

src/cpu/x64/matmul/amx_blocking_heuristics.cpp

@@ -955,9 +955,11 @@ void matmul_amx_blocking_params_micro_t::set_blocking_parameters(
             if (brgemm_k_elems >= K) {
                 k_blk_ = K;
                 k_chunk_size_ = 1;
+                brgemm_batch_size_ = 1;
             } else {
                 k_blk_ = brgemm_k_elems;
                 k_chunk_size_ = 1;
+                brgemm_batch_size_ = 1;
             }
         } else if (current_k_tail == 0
                 && K % (k_blk_ * brgemm_batch_size_) == 0) {

src/cpu/x64/matmul/brgemm_matmul_utils.cpp

@@ -1738,8 +1738,7 @@ void init_aux_values(brgemm_matmul_conf_t &bgmmc,
     bgmmc.buffer_b_chunk_sz = bgmmc.tr_b_dt_sz * rnd_up(bgmmc.N_blk, bgmmc.LDB)
             * rnd_up(bgmmc.K_chunk_elems, bgmmc.wei_k_blk);
 
-    bgmmc.buffer_b_per_thread_sz
-            = bgmmc.buffer_b_chunk_sz * bgmmc.brgemm_batch_size;
+    bgmmc.buffer_b_per_thread_sz = bgmmc.buffer_b_chunk_sz;
 
     bgmmc.buffer_reduce_per_thread_sz = 0;
     if (bgmmc.reduce_kind == matmul_reduce_kind::src) {