cpu: x64: matmul fix wei_k_blk query

Author: densamoilov

src/cpu/x64/matmul/brgemm_matmul.cpp

@@ -1755,7 +1755,7 @@ struct brgemm_matmul_t<isa>::brg_matmul_exec_ctx_t {
 
     dim_t get_data_B_kn_off(int k, int n) const {
         const int wei_k_blk
-                = bgmmc_.is_bf32 ? bgmmc_.orig_wei_k_blk : bgmmc_.wei_k_blk;
+                = bgmmc_.is_bf32 ? get_wei_k_blk(f32) : bgmmc_.wei_k_blk;
         const int k_idx = bgmmc_.blocked_B ? k / wei_k_blk : k;
         const int n_idx = bgmmc_.blocked_B ? n / bgmmc_.wei_n_blk : n;
         const int int4_fac = bgmmc_.is_int4_weights ? 2 : 1;

src/cpu/x64/matmul/brgemm_matmul_utils.cpp

@@ -81,6 +81,16 @@ int get_n_block_from_tag(format_tag_t matrix_b_tag) {
     }
 }
 
+int get_wei_k_blk(data_type_t wei_dt) {
+    // Fixed outer block size.
+    const int k_outer_block = 16;
+
+    // VNNI granularity determines the inner block size along K.
+    const int k_inner_block = data_type_vnni_granularity(wei_dt);
+
+    return k_outer_block * k_inner_block;
+}
+
 void mem_advice_init(brgemm_matmul_conf_t &bgmmc) {
 
     dim_t parallel_work_amount = bgmmc.batch * bgmmc.M_chunks * bgmmc.N_chunks;
@@ -1528,9 +1538,7 @@ status_t init_brgemm_matmul_conf(cpu_isa_t isa, brgemm_matmul_conf_t &bgmmc,
 
     VCONDCHECK_BG(bgmmc.required_k_granularity > 0, VERBOSE_BLOCKING_FAIL, "");
 
-    bgmmc.wei_k_blk = bm_conf_utils.get_wei_k_blk();
-    bgmmc.orig_wei_k_blk
-            = bm_conf_utils.get_wei_k_blk(/*use_orig_wei_dt=*/true);
+    bgmmc.wei_k_blk = get_wei_k_blk(bgmmc.wei_dt);
 
     VCHECK_BG(bm_conf_utils.set_or_check_B_tag(weights_md, helper),
             VERBOSE_UNSUPPORTED_TAG);

src/cpu/x64/matmul/brgemm_matmul_utils.hpp

@@ -94,7 +94,7 @@ struct brgemm_matmul_conf_t {
     dim_t LDA, LDB, LDC, LDD;
     dim_t LDB2;
     int brgemm_batch_size, brgemm_batch_tail_size;
-    int wei_n_blk, wei_k_blk, orig_wei_k_blk;
+    int wei_n_blk, wei_k_blk;
     brgemm_batch_kind_t brg_type;
     bool is_macro_heuristics;
 
@@ -282,32 +282,6 @@ struct brgemm_matmul_conf_utils_t {
                 && check_b_layout_blocked_by_n(bgmmc.wei_tag);
     }
 
-    /**
-     * Returns the total block size along the K dimension, as the product of
-     * the fixed outer block size and the VNNI granularity.
-     *
-     * Example: For format tag BA16a16b4a, the block size is
-     * 16 (outer) * 4 (VNNI granularity) = 64.
-     *
-     * @param use_orig_wei_dt If true, use the original weight data type to
-     *   determine block size. If false, use the compute data type type.
-     * @return The total K dimension block size.
-     */
-    int get_wei_k_blk(bool use_orig_wei_dt = false) const {
-        // No blocking is used in GEMV mode.
-        if (bgmmc.is_gemv) return 1;
-
-        // Fixed outer block size.
-        const int k_outer_block = 16;
-
-        // VNNI granularity determines the inner block size along K.
-        const data_type_t wei_dt
-                = use_orig_wei_dt ? bgmmc.orig_wei_dt : bgmmc.wei_dt;
-        const int k_inner_block = data_type_vnni_granularity(wei_dt);
-
-        return k_outer_block * k_inner_block;
-    }
-
     inline bool use_buffer_b(bool use_heuristic = true) const {
         if (bgmmc.is_runtime_N) return true;
         if (bgmmc.is_bf16_with_int_wei) return true;
@@ -483,6 +457,19 @@ void mem_advice_init(brgemm_matmul_conf_t &bgmmc);
 
 bool is_batch_layout_trivial(const memory_desc_wrapper &mdw, const dim_t batch);
 
+/**
+ * Returns the total block size along the K dimension, as the product of
+ * the fixed outer block size and the VNNI granularity.
+ *
+ * Example: For format tag BA16a16b4a, the block size is
+ * 16 (outer) * 4 (VNNI granularity) = 64.
+ *
+ * @param wei_dt Weights data type.
+ *
+ * @return The total K dimension block size.
+ */
+int get_wei_k_blk(data_type_t wei_dt);
+
 } // namespace matmul
 } // namespace x64
 } // namespace cpu