CopyB with transpose works for u8i8 case

Author: v-Golubev

src/common/snippets/src/pass/explicit_transpose_matmul_inputs.cpp

@@ -76,7 +76,6 @@ ov::snippets::pass::ExplicitTransposeMatMulInputs::ExplicitTransposeMatMulInputs
             ov::pass::pattern::any_input(ov::pass::pattern::has_static_shape()),
             ov::pass::pattern::any_input(ov::pass::pattern::has_static_shape()));
 
-    const bool skip_transpose_b_extraction = std::getenv("TRANSPOSE_B");
     register_matcher(std::make_shared<ov::pass::pattern::Matcher>(m_matmul0, matcher_name),
         [=](ov::pass::pattern::Matcher &m) {
             OV_ITT_SCOPED_TASK(ov::pass::itt::domains::SnippetsTransform, "Snippets::op::ExplicitTransposeMatMulInputs")
@@ -92,6 +91,7 @@ ov::snippets::pass::ExplicitTransposeMatMulInputs::ExplicitTransposeMatMulInputs
                 matmul->set_transpose_a(false);
                 rewritten |= true;
             }
+            const bool skip_transpose_b_extraction = std::getenv("TRANSPOSE_B") && matmul->get_input_element_type(1) != ov::element::f32;
             if (matmul->get_transpose_b() && !skip_transpose_b_extraction) {
                 extract(matmul->input(1));
                 std::cout << "[ INFO ] ExplicitTransposeMatMulInputs is finished for B input\n";

src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.cpp

@@ -22,8 +22,6 @@ using namespace dnnl::impl::cpu::x64;
 namespace ov {
 namespace intel_cpu {
 
-const bool transpose_b_enable = std::getenv("TRANSPOSE_B");
-
 jit_brgemm_copy_b_emitter::jit_brgemm_copy_b_emitter(jit_generator* h, cpu_isa_t isa, const  ov::snippets::lowered::ExpressionPtr& expr)
     : jit_emitter(h, isa) {
     in_out_type_ = emitter_in_out_map::gpr_to_gpr;
@@ -43,16 +41,17 @@ jit_brgemm_copy_b_emitter::jit_brgemm_copy_b_emitter(jit_generator* h, cpu_isa_t
     const auto transposed_shape = layout.empty() ? original_shape : snippets::utils::get_planar_vdims(original_shape, layout);
     const size_t N = transposed_shape.back();
 
-    dnnl_format_tag_t format = dnnl_abcd;
+    m_format = dnnl_abcd;
     size_t wei_stride = 0;
     if (layout == VectorDims{0, 2, 1, 3}) {
-        format = dnnl_acbd;
+        std::cout << "wei stride is needed!!\n";
+        m_format = dnnl_acbd;
         wei_stride = jit_brgemm_emitter::get_in_leading_dim(original_shape, layout);
+    } else if (layout == VectorDims{0, 1, 3, 2}) {
+        std::cout << "transposed copy_b shape\n";
+        m_format = dnnl_abdc;
     }
 
-    if (transpose_b_enable)
-        format = dnnl_abdc;
-
     std::cout << "[ INFO ] CopyBEmitter is processing...\n";
     std::cout << "\tshape = " << ov::PartialShape(original_shape) << std::endl;
     std::cout << "\tlayout = " << ov::PartialShape(layout) << std::endl;
@@ -77,7 +76,7 @@ jit_brgemm_copy_b_emitter::jit_brgemm_copy_b_emitter(jit_generator* h, cpu_isa_t
     OV_CPU_JIT_EMITTER_ASSERT(!one_of(m_brg_weight_etype, element::bf16, element::i8), "doesn't support precision ", m_brg_weight_etype);
     const auto repacking_buffer_shape = brgemm_repack->get_repacking_buffer_shape();
     OV_CPU_JIT_EMITTER_ASSERT(!repacking_buffer_shape.empty(), "Repacking buffer shape mustn't be empty");
-    size_t LDB = transpose_b_enable ? 384 : repacking_buffer_shape.back();
+    size_t LDB = repacking_buffer_shape.back();
     if (auto val = std::getenv("LDB")) {
         LDB = std::atoi(val);
     }
@@ -93,18 +92,17 @@ jit_brgemm_copy_b_emitter::jit_brgemm_copy_b_emitter(jit_generator* h, cpu_isa_t
     const auto src_dt = static_cast<dnnl_data_type_t>(DnnlExtensionUtils::ElementTypeToDataType(brg_src_etype));
     const auto wei_dt = static_cast<dnnl_data_type_t>(DnnlExtensionUtils::ElementTypeToDataType(m_brg_weight_etype));
 
-    init_brgemm_copy(m_kernel, format, wei_stride, N, m_inner_N_block, m_inner_N_tail, LDB, m_K_blk, use_amx, src_dt, wei_dt);
+    init_brgemm_copy(m_kernel, N, m_inner_N_block, m_inner_N_tail, LDB, m_K_blk, use_amx, src_dt, wei_dt, wei_stride);
 }
 
 void jit_brgemm_copy_b_emitter::init_brgemm_copy(std::unique_ptr<matmul::jit_brgemm_matmul_copy_b_t>& kernel,
-                                                 dnnl_format_tag_t format, size_t wei_stride,
                                                  size_t N, size_t N_blk, size_t N_tail, size_t LDB, size_t K,
-                                                 bool is_with_amx, dnnl_data_type_t src_dt, dnnl_data_type_t wei_dt) const {
+                                                 bool is_with_amx, dnnl_data_type_t src_dt, dnnl_data_type_t wei_dt, size_t wei_stride) const {
     matmul::brgemm_matmul_conf_t brgCopyKernelConf;
     brgCopyKernelConf.src_dt = src_dt;
     brgCopyKernelConf.wei_dt = wei_dt;
     brgCopyKernelConf.wei_n_blk = static_cast<int>(N_blk);
-    brgCopyKernelConf.wei_tag = format;
+    brgCopyKernelConf.wei_tag = m_format;
     brgCopyKernelConf.copy_B_wei_stride = wei_stride;
     brgCopyKernelConf.LDB = static_cast<dim_t>(LDB);
     brgCopyKernelConf.N =  static_cast<dim_t>(N);
@@ -148,10 +146,12 @@ void jit_brgemm_copy_b_emitter::emit_impl(const std::vector<size_t>& in, const s
     Xbyak::Reg64 comp(static_cast<int>(m_with_comp ? out[1] : 0));
 
     const size_t data_size = m_brg_weight_etype.size();
+    const size_t K_scale = m_format == dnnl_abdc ? m_K_blk : 1;
     for (size_t nb = 0; nb < div_up(m_N_blk, m_inner_N_block); nb++) {
-        const size_t offset_in = m_in_offset + nb * m_inner_N_block * data_size;
+        const size_t offset_in = m_in_offset + nb * m_inner_N_block * K_scale * data_size;
         const size_t offset_out = m_out_offset + nb * m_inner_N_block * m_brgemmVNNIFactor * data_size;
         const size_t offset_comp = m_with_comp ? m_comp_offset + nb * m_inner_N_block * sizeof(int32_t) : 0;
+        std::cout << "offset in [" << nb << "] = " << offset_in << std::endl;
 
         const bool is_N_tail = (m_N_blk - nb * m_inner_N_block < m_inner_N_block);
         const auto current_N_blk = is_N_tail ? m_inner_N_tail : m_inner_N_block;

src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.hpp

@@ -27,9 +27,8 @@ class jit_brgemm_copy_b_emitter : public jit_emitter {
     void emit_impl(const std::vector<size_t>& in, const std::vector<size_t>& out) const override;
 
     void init_brgemm_copy(std::unique_ptr<dnnl::impl::cpu::x64::matmul::jit_brgemm_matmul_copy_b_t>& kernel,
-                          dnnl_format_tag_t format, size_t wei_stride,
                           size_t N, size_t N_blk, size_t N_tail, size_t LDB, size_t K,
-                          bool is_with_amx, dnnl_data_type_t dt_in0, dnnl_data_type_t dt_in1) const;
+                          bool is_with_amx, dnnl_data_type_t dt_in0, dnnl_data_type_t dt_in1, size_t wei_stride) const;
     void emit_kernel_call(const dnnl::impl::cpu::x64::matmul::jit_brgemm_matmul_copy_b_t* kernel,
                           Xbyak::Reg64 src, Xbyak::Reg64 dst, Xbyak::Reg64 comp, size_t N, size_t K,
                           size_t offset_in, size_t offset_out, size_t offset_comp) const;
@@ -56,6 +55,8 @@ class jit_brgemm_copy_b_emitter : public jit_emitter {
 
     bool m_with_comp = false;
 
+    dnnl_format_tag_t m_format;
+
 #ifdef SNIPPETS_DEBUG_CAPS
     friend std::string init_info_jit_brgemm_copy_b_emitter(const jit_brgemm_copy_b_emitter *emitter);
 #endif