cpu: x64: matmul: fix wei_n_blk when using a b buffer

yair-obodovsky · tprimak · commit 642d18b7fe46 · 2025-09-18T19:03:21.000-10:00
(cherry picked from commit 216a37e)
diff --git a/src/cpu/x64/matmul/brgemm_matmul_utils.cpp b/src/cpu/x64/matmul/brgemm_matmul_utils.cpp
@@ -1541,12 +1541,21 @@ status_t init_brgemm_matmul_conf(cpu_isa_t isa, brgemm_matmul_conf_t &bgmmc,
     VCHECK_BG(compute_blocking_heuristic(bgmmc, bm_conf_utils),
             VERBOSE_BLOCKING_FAIL, "");
 
-    if (bgmmc.wei_n_blk > bgmmc.N_blk
-            && IMPLICATION(
-                    bgmmc.N == bgmmc.N_blk, bgmmc.N >= bgmmc.wei_n_blk)) {
+    if (bgmmc.wei_n_blk > bgmmc.N_blk && bgmmc.N != bgmmc.N_blk) {
         assert(!bgmmc.is_runtime_N
                 && "N_blk should not be adjusted for runtime N");
-        bgmmc.wei_n_blk = bgmmc.N_blk;
+        if (bgmmc.use_buffer_b) {
+            // Copy kernels require the B buffer to be aligned.
+            // ZMM registers are used without masking;
+            // Two YMMs are used in the AVX2 case with the same granularity.
+            size_t n_elements_in_wei_zmm = platform::get_cache_line_size()
+                    / (data_type_vnni_granularity(bgmmc.wei_dt)
+                            * bgmmc.tr_b_dt_sz);
+            bgmmc.wei_n_blk = rnd_up(bgmmc.N_blk, n_elements_in_wei_zmm);
+        } else {
+            bgmmc.wei_n_blk = bgmmc.N_blk;
+        }
+
         VCHECK_BG(bm_conf_utils.update_and_check_B_tag(
                           weights_md, bgmmc.wei_n_blk, helper),
                 VERBOSE_UNSUPPORTED_TAG);
