diff --git a/src/common/snippets/src/lowered/pass/cleanup_loop_offsets.cpp b/src/common/snippets/src/lowered/pass/cleanup_loop_offsets.cpp index 5dc6ad1e136..2b5162d8972 100644 --- a/src/common/snippets/src/lowered/pass/cleanup_loop_offsets.cpp +++ b/src/common/snippets/src/lowered/pass/cleanup_loop_offsets.cpp @@ -31,13 +31,15 @@ bool CleanupLoopOffsets::run(lowered::LinearIR& linear_ir, lowered::LinearIR::co } if (auto outer_loop_end = as_type_ptr(next_node)) { const auto& is_incremented = loop_end->get_is_incremented(); + const auto& data_sizes = loop_end->get_element_type_sizes(); auto fin_offsets = loop_end->get_finalization_offsets(); std::unordered_map per_port_connector_offset; const auto& loop_inputs = expr_it->get()->get_input_port_connectors(); for (size_t i = 0; i < fin_offsets.size(); i++) per_port_connector_offset[loop_inputs[i]] = i; - const auto outer_is_incremented = outer_loop_end->get_is_incremented(); + const auto& outer_is_incremented = outer_loop_end->get_is_incremented(); + const auto& outer_data_sizes = outer_loop_end->get_element_type_sizes(); const auto outer_increment = static_cast(outer_loop_end->get_increment()); auto outer_ptr_increments = outer_loop_end->get_ptr_increments(); const auto& outer_loop_inputs = next_expr_it->get()->get_input_port_connectors(); @@ -47,7 +49,7 @@ bool CleanupLoopOffsets::run(lowered::LinearIR& linear_ir, lowered::LinearIR::co const auto& managed_connector = outer_loop_inputs[i]; const auto& found = per_port_connector_offset.find(managed_connector); if (found != per_port_connector_offset.end()) { - if (!is_incremented[found->second]) + if (!is_incremented[found->second] || outer_data_sizes[i] != data_sizes[found->second]) continue; // Since data ptr is incremented on [ptr_increment x increment], // we should guarantee proportionality of ptr shifts. diff --git a/src/common/snippets/src/lowered/pass/fuse_loops.cpp b/src/common/snippets/src/lowered/pass/fuse_loops.cpp index 5944d4f41ea..035e5f37dc3 100644 --- a/src/common/snippets/src/lowered/pass/fuse_loops.cpp +++ b/src/common/snippets/src/lowered/pass/fuse_loops.cpp @@ -248,7 +248,7 @@ bool FuseLoops::run(LinearIR& linear_ir, lowered::LinearIR::constExprIt begin, l const auto upper_loop_id = upper_loop_ids[loop_idx]; OPENVINO_ASSERT(current_loop_id != upper_loop_id, - "Loops cannot have parents of entry points with the same identifier"); + "Loops cannot have parents of entry points with the same identifier (", upper_loop_id, ")"); if (fuse_upper_into_current(linear_ir, loop_manager, entry_point.expr_port, current_loop_id, upper_loop_id, current_loop_begin_pos, current_loop_end_pos)) { was_fusion_up = true; diff --git a/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.cpp b/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.cpp index 95e73ba9175..2840b6abbda 100644 --- a/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.cpp +++ b/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.cpp @@ -30,9 +30,6 @@ jit_brgemm_copy_b_emitter::jit_brgemm_copy_b_emitter(jit_generator* h, cpu_isa_t if (!brgemm_repack) OV_CPU_JIT_EMITTER_THROW("expects BrgemmCopyB node"); - m_brgemm_prc_in0 = brgemm_repack->get_src_element_type(); - m_brgemm_prc_in1 = brgemm_repack->get_input_element_type(0); - m_brgemmVNNIFactor = 4 / m_brgemm_prc_in0.size(); m_with_comp = brgemm_repack->is_with_compensations(); m_in_offset = brgemm_repack->get_offset_in(); m_out_offset = brgemm_repack->get_offset_out(); @@ -49,36 +46,52 @@ jit_brgemm_copy_b_emitter::jit_brgemm_copy_b_emitter(jit_generator* h, cpu_isa_t leading_dimension = jit_brgemm_emitter::get_in_leading_dim(original_shape, layout); } - m_N = *(transposed_shape.rbegin()); - m_K = *(transposed_shape.rbegin() + 1); + const auto& in_subtensor = in_desc->get_subtensor(); + m_N_blk = *in_subtensor.rbegin(); + m_K_blk = *++in_subtensor.rbegin(); + OV_CPU_JIT_EMITTER_ASSERT(m_N_blk <= *transposed_shape.rbegin() && m_K_blk <= *++transposed_shape.rbegin(), + "BrgemmCopyB has incompatible subtensor dimensions"); + m_inner_N_block = brgemm_repack->get_n_inner_block_size(); + m_inner_N_tail = m_N_blk % m_inner_N_block; - m_N_blk = brgemm_repack->get_n_block_size(); - m_K_blk = brgemm_repack->get_k_block_size(); + OV_CPU_JIT_EMITTER_ASSERT(expr->get_output_port_descriptor(0)->get_subtensor() == in_subtensor, "output and input subtensors must be equal"); + if (m_with_comp) { + const auto& compensations_subtensor = expr->get_output_port_descriptor(1)->get_subtensor(); + OV_CPU_JIT_EMITTER_ASSERT( + *compensations_subtensor.rbegin() == m_N_blk && *++compensations_subtensor.rbegin() == 1, + "compensations subtensor must be {1, m_N_blk}"); + } - m_N_tail = m_N % m_N_blk; - m_K_tail = m_K % m_K_blk; - m_LDB = m_brgemm_prc_in1 == ov::element::f32 ? leading_dimension : rnd_up(m_N, m_N_blk); + 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"); + const auto& LDB = repacking_buffer_shape.back(); - const auto dt_in0 = static_cast(DnnlExtensionUtils::ElementTypeToDataType(m_brgemm_prc_in0)); - const auto dt_in1 = static_cast(DnnlExtensionUtils::ElementTypeToDataType(m_brgemm_prc_in1)); + const auto& brg_src_etype = brgemm_repack->get_src_element_type(); + m_brg_weight_etype = brgemm_repack->get_input_element_type(0); + m_brgemmVNNIFactor = brgemm_repack->get_brgemm_vnni_factor(); - const bool isAMXSupported = mayiuse(avx512_core_amx); - const auto use_amx = isAMXSupported && m_brgemm_prc_in0 != ov::element::f32 && (m_K % m_brgemmVNNIFactor == 0) && (m_N % m_brgemmVNNIFactor == 0); - init_brgemm_copy(m_kernel, leading_dimension, m_N_blk, m_N_tail, m_LDB, m_K - m_K_tail, use_amx, dt_in0, dt_in1); + const auto use_amx = mayiuse(avx512_core_amx) && brg_src_etype != ov::element::f32 && + (m_K_blk % m_brgemmVNNIFactor == 0) && (m_N_blk % m_brgemmVNNIFactor == 0); + + const auto src_dt = static_cast(DnnlExtensionUtils::ElementTypeToDataType(brg_src_etype)); + const auto wei_dt = static_cast(DnnlExtensionUtils::ElementTypeToDataType(m_brg_weight_etype)); + + init_brgemm_copy(m_kernel, leading_dimension, m_inner_N_block, m_inner_N_tail, LDB, m_K_blk, use_amx, src_dt, wei_dt); } void jit_brgemm_copy_b_emitter::init_brgemm_copy(std::unique_ptr& kernel, - 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 { + 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 { matmul::brgemm_matmul_conf_t brgCopyKernelConf; - brgCopyKernelConf.src_dt = dt_in0; - brgCopyKernelConf.wei_dt = dt_in1; + brgCopyKernelConf.src_dt = src_dt; + brgCopyKernelConf.wei_dt = wei_dt; brgCopyKernelConf.wei_n_blk = static_cast(N_blk); brgCopyKernelConf.wei_tag = dnnl_abcd; // What's about other ranks? brgCopyKernelConf.copy_B_wei_stride = 0; brgCopyKernelConf.LDB = static_cast(LDB); brgCopyKernelConf.N = static_cast(N); - brgCopyKernelConf.N_tail = static_cast(N_tail); + brgCopyKernelConf.N_tail = static_cast(N_tail); brgCopyKernelConf.N_blk = static_cast(N_blk); brgCopyKernelConf.K = static_cast(K); brgCopyKernelConf.K_blk = static_cast(K); @@ -91,8 +104,8 @@ void jit_brgemm_copy_b_emitter::init_brgemm_copy(std::unique_ptr& in, - const std::vector& out) const { - if (host_isa_ == cpu::x64::avx512_core) { - Xbyak::Reg64 src(static_cast(in[0])); - Xbyak::Reg64 dst(static_cast(out[0])); - Xbyak::Reg64 comp(static_cast(0)); // Compensations. Default reg idx is 0 if there aren't the compensations - if (m_with_comp) { - if (out.size() != 2) { - OV_CPU_JIT_EMITTER_THROW("with compensations requires separate register for them"); - } - comp = Xbyak::Reg64(static_cast(out[1])); - } +void jit_brgemm_copy_b_emitter::validate_arguments(const std::vector &in, const std::vector &out) const { + OV_CPU_JIT_EMITTER_ASSERT(in.size() == 1, "expects 1 input"); + OV_CPU_JIT_EMITTER_ASSERT((m_with_comp && out.size() == 2) || (!m_with_comp && out.size() == 1), + "expects 2 outputs if there are compensations"); +} - const size_t data_size = m_brgemm_prc_in1.size(); - for (size_t nb = 0; nb < div_up(m_N, m_N_blk); nb++) { - const size_t offset_in = m_in_offset + nb * m_N_blk * data_size; - const size_t offset_out = m_out_offset + nb * m_N_blk * m_brgemmVNNIFactor * data_size; - const size_t offset_comp = m_with_comp ? m_comp_offset + nb * m_N_blk * sizeof(int32_t) : 0; +void jit_brgemm_copy_b_emitter::emit_impl(const std::vector& in, const std::vector& out) const { + validate_arguments(in, out); + OV_CPU_JIT_EMITTER_ASSERT(host_isa_ == cpu::x64::avx512_core, "requires at least avx512_core instruction set"); - const bool is_N_tail = (m_N - nb * m_N_blk < m_N_blk); - const auto current_N_blk = is_N_tail ? m_N_tail : m_N_blk; + Xbyak::Reg64 src(static_cast(in[0])); + Xbyak::Reg64 dst(static_cast(out[0])); + Xbyak::Reg64 comp(static_cast(m_with_comp ? out[1] : 0)); - emit_kernel_call(m_kernel.get(), src, dst, comp, current_N_blk, m_K, offset_in, offset_out, offset_comp); - } - } else { - OV_CPU_JIT_EMITTER_THROW("requires at least avx512_core instruction set"); + const size_t data_size = m_brg_weight_etype.size(); + 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_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; + + 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; + + emit_kernel_call(m_kernel.get(), src, dst, comp, current_N_blk, m_K_blk, offset_in, offset_out, offset_comp); } } void jit_brgemm_copy_b_emitter::emit_kernel_call(const matmul::jit_brgemm_matmul_copy_b_t* kernel, Reg64 src, Reg64 dst, Reg64 comp, - size_t N, size_t K, size_t offset_in, size_t offset_out, size_t offset_comp) const { + size_t N, size_t K, size_t offset_in, size_t offset_out, size_t offset_comp) const { const auto data_ptr = [&](Xmm xmm, Xbyak::Reg64 reg, size_t bytes_offset) { h->uni_vmovq(reg, xmm); if (bytes_offset) h->add(reg, bytes_offset); @@ -199,11 +209,12 @@ void jit_brgemm_copy_b_emitter::emit_kernel_call(const matmul::jit_brgemm_matmul internal_call_postamble(); } -void jit_brgemm_copy_b_emitter::execute(matmul::jit_brgemm_matmul_copy_b_t *kernel, const void *src, - const void *dst, const void *comp, size_t N, size_t K) { - if (!kernel) - OV_CPU_JIT_EMITTER_THROW("Kernel hasn't been created"); - +void jit_brgemm_copy_b_emitter::execute(matmul::jit_brgemm_matmul_copy_b_t* kernel, + const void* src, + const void* dst, + const void* comp, + size_t N, + size_t K) { auto ctx = dnnl::impl::cpu::x64::matmul::jit_brgemm_matmul_copy_b_t::ctx_t(); ctx.current_N_blk = N; ctx.src = src; @@ -214,6 +225,7 @@ void jit_brgemm_copy_b_emitter::execute(matmul::jit_brgemm_matmul_copy_b_t *kern ctx.current_K_start = 0; ctx.current_K_iters = K; + OV_CPU_JIT_EMITTER_ASSERT(kernel, "Kernel hasn't been created"); (*kernel)(&ctx); } diff --git a/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.hpp b/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.hpp index f11c1e84c29..00cf922a905 100644 --- a/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.hpp +++ b/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.hpp @@ -23,6 +23,7 @@ public: } private: + void validate_arguments(const std::vector &in, const std::vector &out) const override; void emit_impl(const std::vector& in, const std::vector& out) const override; void init_brgemm_copy(std::unique_ptr& kernel, @@ -36,18 +37,24 @@ private: const void* src, const void* dst, const void* comp, size_t N, size_t K); std::unique_ptr m_kernel; + ov::element::Type m_brg_weight_etype; - ov::element::Type m_brgemm_prc_in0, m_brgemm_prc_in1; - size_t m_N, m_N_blk, m_N_tail; - size_t m_K, m_K_blk, m_K_tail; - size_t m_LDB; - size_t m_brgemmVNNIFactor; - bool m_with_comp = false; + // Block size which is set by snippets: it is usually shared between brgemm and brgemm_copy_b nodes + size_t m_N_blk = 0lu; + // Block size which is used by the internal OneDNN implementation. + // It is used in snippets emitter to iterate through input/output data and call OneDNN kernel + size_t m_inner_N_block = 0lu; + size_t m_inner_N_tail = 0lu; + + size_t m_K_blk = 0lu; + size_t m_brgemmVNNIFactor = 0lu; size_t m_in_offset = 0lu; size_t m_out_offset = 0lu; size_t m_comp_offset = 0lu; + bool m_with_comp = false; + #ifdef SNIPPETS_DEBUG_CAPS friend std::string init_info_jit_brgemm_copy_b_emitter(const jit_brgemm_copy_b_emitter *emitter); #endif diff --git a/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_emitter.cpp b/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_emitter.cpp index 1cf1c528144..78bddccdffa 100644 --- a/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_emitter.cpp +++ b/src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_emitter.cpp @@ -78,9 +78,9 @@ jit_brgemm_emitter::jit_brgemm_emitter(jit_generator* h, cpu_isa_t isa, const ov init_in_scheduling_params(input_0_desc); if (brgemm_node->is_with_data_repacking()) { - const auto& brgemm_copy = brgemm_node->get_brgemm_copy(); - const auto& allocated_shape = brgemm_copy->get_data_repacking_shape(input_1_desc->get_shape()); - leading_dimensions.push_back(*allocated_shape.rbegin()); + const auto repacking_buffer_shape = brgemm_node->get_brgemm_copy()->get_repacking_buffer_shape(); + OV_CPU_JIT_EMITTER_ASSERT(!repacking_buffer_shape.empty(), "Repacking buffer shape mustn't be empty"); + leading_dimensions.push_back(repacking_buffer_shape.back()); } else { init_in_scheduling_params(input_1_desc); } @@ -89,7 +89,6 @@ jit_brgemm_emitter::jit_brgemm_emitter(jit_generator* h, cpu_isa_t isa, const ov const auto& brg0Prc = brgemm_node->get_input_element_type(0); const auto& brg1Prc = brgemm_node->get_input_element_type(1); - m_with_comp = brgemm_node->is_with_compensations(); m_with_scratch = brgemm_node->is_with_scratchpad(); const auto& output_subtensor = output_desc->get_subtensor(); @@ -113,6 +112,7 @@ jit_brgemm_emitter::jit_brgemm_emitter(jit_generator* h, cpu_isa_t isa, const ov m_ctx.dt_in1 = static_cast(DnnlExtensionUtils::ElementTypeToDataType(brg1Prc)); m_ctx.beta = brgemm_node->get_beta(); m_ctx.is_with_amx = brgemm_node->is_amx(); + m_ctx.is_with_comp = brgemm_node->is_with_compensations(); init_brgemm_kernel(m_ctx, m_kernel); @@ -155,8 +155,6 @@ void jit_brgemm_emitter::init_brgemm_kernel(brgemmCtx& ctx, std::unique_ptr m_kernel = nullptr; bool m_with_scratch = false; - bool m_with_comp = false; size_t m_load_offset_a = 0lu; size_t m_load_offset_b = 0lu; diff --git a/src/plugins/intel_cpu/src/emitters/snippets/x64/verbose.cpp b/src/plugins/intel_cpu/src/emitters/snippets/x64/verbose.cpp index e1e898b03e6..5c6c9bca408 100644 --- a/src/plugins/intel_cpu/src/emitters/snippets/x64/verbose.cpp +++ b/src/plugins/intel_cpu/src/emitters/snippets/x64/verbose.cpp @@ -119,8 +119,7 @@ std::string init_info_jit_brgemm_emitter(const jit_brgemm_emitter *emitter) { << " m_load_offset_b:" << emitter->m_load_offset_b << " m_load_offset_scratch:" << emitter->m_load_offset_scratch << " m_store_offset_c:" << emitter->m_store_offset_c - << " m_with_scratch:" << emitter->m_with_scratch - << " m_with_comp:" << emitter->m_with_comp; + << " m_with_scratch:" << emitter->m_with_scratch; return ss.str(); } @@ -128,20 +127,16 @@ std::string init_info_jit_brgemm_emitter(const jit_brgemm_emitter *emitter) { std::string init_info_jit_brgemm_copy_b_emitter(const jit_brgemm_copy_b_emitter *emitter) { std::stringstream ss; ss << "Emitter_type_name:jit_brgemm_copy_b_emitter" - << " m_LDB:" << emitter->m_LDB - << " m_K:" << emitter->m_K - << " m_K_blk:" << emitter->m_K_blk - << " m_K_tail:" << emitter->m_K_tail - << " m_N:" << emitter->m_N + << " m_brg_weight_etype:" << emitter->m_brg_weight_etype << " m_N_blk:" << emitter->m_N_blk - << " m_N_tail:" << emitter->m_N_tail - << " m_brgemm_prc_in0:" << emitter->m_brgemm_prc_in0 - << " m_brgemm_prc_in1:" << emitter->m_brgemm_prc_in1 + << " m_inner_N_block:" << emitter->m_inner_N_block + << " m_inner_N_tail:" << emitter->m_inner_N_tail + << " m_K_blk:" << emitter->m_K_blk << " m_brgemmVNNIFactor:" << emitter->m_brgemmVNNIFactor - << " m_with_comp:" << emitter->m_with_comp << " m_in_offset:" << emitter->m_in_offset << " m_out_offset:" << emitter->m_out_offset - << ",m_comp_offset:" << emitter->m_comp_offset; + << " m_comp_offset:" << emitter->m_comp_offset + << " m_with_comp:" << emitter->m_with_comp; return ss.str(); } diff --git a/src/plugins/intel_cpu/src/transformations/snippets/x64/op/brgemm_copy_b.cpp b/src/plugins/intel_cpu/src/transformations/snippets/x64/op/brgemm_copy_b.cpp index 9bf6d7a1d62..fbf0e01900b 100644 --- a/src/plugins/intel_cpu/src/transformations/snippets/x64/op/brgemm_copy_b.cpp +++ b/src/plugins/intel_cpu/src/transformations/snippets/x64/op/brgemm_copy_b.cpp @@ -18,7 +18,6 @@ intel_cpu::BrgemmCopyB::BrgemmCopyB(const Output& x, const element::Type s std::vector layout_input, const size_t blk_size_k, const size_t blk_size_n) : snippets::modifier::MemoryAccess(1, type == Type::WithCompensations ? 2 : 1), op::Op({x}), m_type(type), m_src_type(src_type) { - m_brgemmVNNIFactor = 4 / m_src_type.size(); set_output_size(type == Type::WithCompensations ? 2 : 1); set_input_port_descriptor({0, offset_in}, 0); set_output_port_descriptor({0, offset_out0}, 0); @@ -27,6 +26,8 @@ intel_cpu::BrgemmCopyB::BrgemmCopyB(const Output& x, const element::Type s } compute_block_size_values(blk_size_k, blk_size_n); custom_constructor_validate_and_infer_types(std::move(layout_input)); + m_brgemmVNNIFactor = 4 / m_src_type.size(); + m_inner_n_block = 16 * m_brgemmVNNIFactor; } intel_cpu::BrgemmCopyB::BrgemmCopyB(const Output& x, const element::Type src_type, const Type type, @@ -34,7 +35,6 @@ intel_cpu::BrgemmCopyB::BrgemmCopyB(const Output& x, const element::Type s std::vector layout_input, const size_t blk_size_k, const size_t blk_size_n) : snippets::modifier::MemoryAccess(1, type == Type::WithCompensations ? 2 : 1), op::Op({x}), m_type(type), m_src_type(src_type) { - m_brgemmVNNIFactor = 4 / m_src_type.size(); set_output_size(type == Type::WithCompensations ? 2 : 1); set_input_port_descriptor(desc_in0, 0); set_output_port_descriptor(desc_out0, 0); @@ -43,6 +43,8 @@ intel_cpu::BrgemmCopyB::BrgemmCopyB(const Output& x, const element::Type s } compute_block_size_values(blk_size_k, blk_size_n); custom_constructor_validate_and_infer_types(std::move(layout_input)); + m_brgemmVNNIFactor = 4 / m_src_type.size(); + m_inner_n_block = 16 * m_brgemmVNNIFactor; } bool BrgemmCopyB::visit_attributes(AttributeVisitor& visitor) { @@ -92,15 +94,18 @@ void intel_cpu::BrgemmCopyB::compute_block_size_values(const size_t blk_size_k, m_N_blk = blk_size_n != 0 ? blk_size_n : *input_shape.rbegin(); } -ov::Shape intel_cpu::BrgemmCopyB::get_data_repacking_shape(const ov::snippets::VectorDims& planar_dims) const { - const auto& N = *planar_dims.rbegin(); - const auto& K = *(planar_dims.rbegin() + 1); - return ov::Shape{rnd_up(K, m_brgemmVNNIFactor), rnd_up(N, m_N_blk)}; +ov::Shape intel_cpu::BrgemmCopyB::get_repacking_buffer_shape() const { + // OneDNN implementation repacks data and always writes the result by m_brgemmVNNIFactor * m_inner_n_block blocks + // despite the actual size of the input data. Because of that we have to round-up the allocation shape to always have enough memory allocated. + // For the details, please see 'copy_4x64' and 'copy_2x32' implementations and usage in onednn/src/cpu/x64/matmul/brgemm_matmul_copy_utils.cpp + return ov::Shape{rnd_up(m_K_blk, m_brgemmVNNIFactor), rnd_up(m_N_blk, m_inner_n_block)}; } -ov::Shape intel_cpu::BrgemmCopyB::get_compensation_shape(const ov::snippets::VectorDims& planar_dims) const { - const auto& N = *planar_dims.rbegin(); - return ov::Shape{rnd_up(N, m_N_blk)}; +ov::Shape intel_cpu::BrgemmCopyB::get_compensations_buffer_shape() const { + // Compensations are computed during repacking, so we need to round-up allocation shape according to m_inner_n_block + // because of OneDNN implementation nuances (as in get_repacking_buffer_shape). + // However, the compensations are computed by N dimension, so K dimension doesn't affect the compensations buffer + return ov::Shape{rnd_up(m_N_blk, m_inner_n_block)}; } std::shared_ptr intel_cpu::BrgemmCopyB::clone_with_new_inputs(const OutputVector& new_args) const { diff --git a/src/plugins/intel_cpu/src/transformations/snippets/x64/op/brgemm_copy_b.hpp b/src/plugins/intel_cpu/src/transformations/snippets/x64/op/brgemm_copy_b.hpp index 3da69e0d889..4df17e71cb2 100644 --- a/src/plugins/intel_cpu/src/transformations/snippets/x64/op/brgemm_copy_b.hpp +++ b/src/plugins/intel_cpu/src/transformations/snippets/x64/op/brgemm_copy_b.hpp @@ -41,11 +41,12 @@ public: size_t get_k_block_size() const { return m_K_blk; } size_t get_n_block_size() const { return m_N_blk; } + size_t get_n_inner_block_size() const { return m_inner_n_block; } void set_k_block_size(size_t block_size) { m_K_blk = block_size; } void set_n_block_size(size_t block_size) { m_N_blk = block_size; } - ov::Shape get_data_repacking_shape(const ov::snippets::VectorDims& planar_dims) const; - ov::Shape get_compensation_shape(const ov::snippets::VectorDims& planar_dims) const; + ov::Shape get_repacking_buffer_shape() const; + ov::Shape get_compensations_buffer_shape() const; Type get_type() const { return m_type; } size_t get_brgemm_vnni_factor() const { return m_brgemmVNNIFactor; } @@ -75,6 +76,9 @@ private: size_t m_K_blk = 0; size_t m_N_blk = 0; + // OneDNN implementation requirement: BrgemmCopyB oneDNN implementation repacks data by m_brgemmVNNIFactor * m_inner_n_block blocks. + // Consequently, in snippets emitter, we need to invoke the oneDNN kernel iterating accordingly to this block + size_t m_inner_n_block = 0; size_t m_brgemmVNNIFactor = 1; }; diff --git a/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.cpp b/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.cpp index 1fb20d6b70e..cc6b769b52d 100644 --- a/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.cpp +++ b/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.cpp @@ -14,18 +14,18 @@ #include "transformations/snippets/x64/op/brgemm_cpu.hpp" #include "transformations/tpp/x64/op/brgemm.hpp" + namespace ov { namespace intel_cpu { namespace pass { using LinearIR = snippets::lowered::LinearIR; using LoopPort = snippets::lowered::LoopPort; -using LoopInfo = snippets::lowered::LoopInfo; using ExpressionPtr = ov::snippets::lowered::ExpressionPtr; -using namespace ov::snippets::lowered::pass; +using namespace ov::snippets::lowered; BrgemmBlocking::BrgemmBlocking() : RangedPass() {} -void BrgemmBlocking::move_new_memory_buffer(snippets::lowered::LinearIR& linear_ir, const snippets::lowered::LinearIR::constExprIt& brgemm_it) { +LinearIR::constExprIt BrgemmBlocking::move_new_memory_buffer(LinearIR& linear_ir, const LinearIR::constExprIt& brgemm_it) { const auto& brgemm_expr = brgemm_it->get(); const auto wsp_expr = brgemm_expr->get_input_port_connector(2)->get_source().get_expr(); const auto wsp_buffer = ov::as_type_ptr(wsp_expr->get_node()); @@ -35,6 +35,25 @@ void BrgemmBlocking::move_new_memory_buffer(snippets::lowered::LinearIR& linear_ const auto wsp_it = linear_ir.find(wsp_expr); linear_ir.move(wsp_it, brgemm_it); } + return std::prev(brgemm_it); +} + +LinearIR::constExprIt BrgemmBlocking::get_loop_begin_pos(LinearIR& linear_ir, const LinearIR::constExprIt& brgemm_it) { + auto loop_begin_it = brgemm_it; + const auto& brgemm_expr = *brgemm_it; + const auto node = brgemm_expr->get_node(); + const auto brgemm = ov::as_type_ptr(node); + const auto brgemm_cpu = ov::as_type_ptr(node); + OPENVINO_ASSERT(brgemm, "get_loop_begin_pos must be called only for Brgemm expression"); + if (brgemm_cpu && brgemm_cpu->is_amx()) { + loop_begin_it = move_new_memory_buffer(linear_ir, brgemm_it); + } + if (brgemm_cpu && brgemm_cpu->is_with_data_repacking()) { + const auto& copy_b = brgemm_cpu->get_brgemm_copy(); + const auto& copy_b_expr = linear_ir.get_expr_by_node(copy_b); + loop_begin_it = linear_ir.find(copy_b_expr); + } + return loop_begin_it; } bool BrgemmBlocking::run(LinearIR& linear_ir, LinearIR::constExprIt begin, LinearIR::constExprIt end) { @@ -79,95 +98,89 @@ bool BrgemmBlocking::run(LinearIR& linear_ir, LinearIR::constExprIt begin, Linea auto in_1_subtensor = in_1_desc->get_subtensor(); auto out_subtensor = out_desc->get_subtensor(); - auto apply_m_blocking = [&]() { - const auto& m = *(out_preordered_dims.rbegin() + 1); - const auto block_size_m = brgemm->get_m_block_size(); - if (block_size_m >= m) { - *(in_0_subtensor.rbegin() + 1) = m; - *(out_subtensor.rbegin() + 1) = m; - return; - } + const auto& m = *++out_preordered_dims.rbegin(); + const auto& n = *out_preordered_dims.rbegin(); + const auto& k = *in_0_planar_dims.rbegin(); + OPENVINO_ASSERT(k == *++in_1_planar_dims.rbegin(), "Brgemm input descriptors have different K dimension value."); - *(in_0_subtensor.rbegin() + 1) = block_size_m; - *(out_subtensor.rbegin() + 1) = block_size_m; - auto loop_begin_it = expr_it, loop_end_it = std::next(expr_it); - std::vector entries{LoopPort(brgemm_expr->get_input_port(0), true), - LoopPort(brgemm_expr->get_input_port(1), false)}; - if (brgemm_cpu) { - if (brgemm_cpu->is_with_compensations()) { - entries.emplace_back(brgemm_expr->get_input_port(2), false); - } else if (brgemm_cpu->is_amx()) { - move_new_memory_buffer(linear_ir, expr_it); - loop_begin_it = std::prev(expr_it); - } - } - std::vector exits{LoopPort(brgemm_expr->get_output_port(0), true)}; - loop_manager->mark_loop(loop_begin_it, loop_end_it, m, block_size_m, 1, entries, exits); - }; + const auto block_size_m = snippets::utils::is_dynamic_value(m) ? brgemm->get_m_block_size() : std::min(brgemm->get_m_block_size(), m); + const auto block_size_n = snippets::utils::is_dynamic_value(n) ? brgemm->get_n_block_size() : std::min(brgemm->get_n_block_size(), n); + const auto block_size_k = snippets::utils::is_dynamic_value(k) ? brgemm->get_k_block_size() : std::min(brgemm->get_k_block_size(), k); - auto apply_n_blocking = [&]() { - const auto& n = *out_preordered_dims.rbegin(); - const auto block_size_n = brgemm->get_n_block_size(); - if (block_size_n >= n) { - *in_1_subtensor.rbegin() = n; - *out_subtensor.rbegin() = n; - return; - } - - *in_1_subtensor.rbegin() = block_size_n; - *out_subtensor.rbegin() = block_size_n; - auto loop_begin_it = expr_it, loop_end_it = std::next(expr_it); - std::vector entries{LoopPort(brgemm_expr->get_input_port(0), false), - LoopPort(brgemm_expr->get_input_port(1), true)}; - if (brgemm_cpu) { - if (brgemm_cpu->is_with_compensations()) { - entries.emplace_back(brgemm_expr->get_input_port(2), true); - } else if (brgemm_cpu->is_amx()) { - move_new_memory_buffer(linear_ir, expr_it); - loop_begin_it = std::prev(expr_it); - } - } - std::vector exits{LoopPort(brgemm_expr->get_output_port(0), true)}; - loop_manager->mark_loop(loop_begin_it, loop_end_it, n, block_size_n, 0, entries, exits); - }; - - auto apply_k_blocking = [&]() { - const auto& k = *in_0_planar_dims.rbegin(); - OPENVINO_ASSERT(k == *(in_1_planar_dims.rbegin() + 1), "Brgemm input descriptors have different K dimension value."); - const auto block_size_k = brgemm->get_k_block_size(); - if (block_size_k >= k) { - *in_0_subtensor.rbegin() = k; - *(in_1_subtensor.rbegin() + 1) = k; - brgemm->set_beta(0.f); - return; - } - - *in_0_subtensor.rbegin() = block_size_k; - *(in_1_subtensor.rbegin() + 1) = block_size_k; - auto loop_begin_it = expr_it, loop_end_it = std::next(expr_it); - std::vector entries{LoopPort(brgemm_expr->get_input_port(0), true, 0), - LoopPort(brgemm_expr->get_input_port(1), true, 1)}; - if (brgemm_cpu) { - if (brgemm_cpu->is_with_compensations()) { - entries.emplace_back(brgemm_expr->get_input_port(2), false, 1); - } else if (brgemm_cpu->is_amx()) { - move_new_memory_buffer(linear_ir, expr_it); - loop_begin_it = std::prev(expr_it); - } - } - std::vector exits{LoopPort(brgemm_expr->get_output_port(0), false)}; - const auto id = loop_manager->mark_loop(loop_begin_it, loop_end_it, k, block_size_k, entries, exits); - const auto loop_info = loop_manager->get_loop_info(id); - loop_info->register_handler(0.f); - }; - - apply_k_blocking(); - apply_n_blocking(); - apply_m_blocking(); + *++in_0_subtensor.rbegin() = block_size_m; + *++out_subtensor.rbegin() = block_size_m; + *in_1_subtensor.rbegin() = block_size_n; + *out_subtensor.rbegin() = block_size_n; + *in_0_subtensor.rbegin() = block_size_k; + *++in_1_subtensor.rbegin() = block_size_k; brgemm_expr->get_input_port_descriptor(0)->set_subtensor(in_0_subtensor); brgemm_expr->get_input_port_descriptor(1)->set_subtensor(in_1_subtensor); brgemm_expr->get_output_port_descriptor(0)->set_subtensor(out_subtensor); + + ov::snippets::lowered::ExpressionPtr copy_b_expr = nullptr; + if (brgemm_cpu && brgemm_cpu->is_with_data_repacking()) { + const auto copy_b = brgemm_cpu->get_brgemm_copy(); + copy_b_expr = linear_ir.get_expr_by_node(copy_b); + + auto data_repacking_subtensor = copy_b_expr->get_input_port_descriptor(0)->get_subtensor(); + *data_repacking_subtensor.rbegin() = block_size_n; + *++data_repacking_subtensor.rbegin() = block_size_k; + + copy_b_expr->get_input_port_descriptor(0)->set_subtensor(data_repacking_subtensor); + copy_b_expr->get_output_port_descriptor(0)->set_subtensor(data_repacking_subtensor); + if (copy_b->is_with_compensations()) { + auto compensations_subtensor = copy_b_expr->get_output_port_descriptor(1)->get_subtensor(); + // Compensations are computed by N dimension + *compensations_subtensor.rbegin() = block_size_n; + *++compensations_subtensor.rbegin() = 1; + copy_b_expr->get_output_port_descriptor(1)->set_subtensor(compensations_subtensor); + } + } + + auto mark_m_blocking = [&]() { + const auto loop_begin_it = get_loop_begin_pos(linear_ir, expr_it); + const auto loop_end_it = std::next(expr_it); + + const std::vector entries{ + LoopPort(brgemm_expr->get_input_port(0), true), + LoopPort(brgemm_cpu && brgemm_cpu->is_with_data_repacking() ? copy_b_expr->get_input_port(0) : brgemm_expr->get_input_port(1), false)}; + const std::vector exits{LoopPort(brgemm_expr->get_output_port(0), true)}; + loop_manager->mark_loop(loop_begin_it, loop_end_it, m, block_size_m, 1, entries, exits); + }; + + auto mark_n_blocking = [&]() { + const auto loop_begin_it = get_loop_begin_pos(linear_ir, expr_it); + const auto loop_end_it = std::next(expr_it); + + const std::vector entries{ + LoopPort(brgemm_expr->get_input_port(0), false), + LoopPort(brgemm_cpu && brgemm_cpu->is_with_data_repacking() ? copy_b_expr->get_input_port(0) : brgemm_expr->get_input_port(1), true)}; + const std::vector exits{LoopPort(brgemm_expr->get_output_port(0), true)}; + loop_manager->mark_loop(loop_begin_it, loop_end_it, n, block_size_n, 0, entries, exits); + }; + + auto mark_k_blocking = [&]() { + const auto loop_begin_it = get_loop_begin_pos(linear_ir, expr_it); + const auto loop_end_it = std::next(expr_it); + + const std::vector entries{ + LoopPort(brgemm_expr->get_input_port(0), true, 0), + LoopPort(brgemm_cpu && brgemm_cpu->is_with_data_repacking() ? copy_b_expr->get_input_port(0) : brgemm_expr->get_input_port(1), true, 1)}; + const std::vector exits{LoopPort(brgemm_expr->get_output_port(0), false)}; + const auto id = loop_manager->mark_loop(loop_begin_it, loop_end_it, k, block_size_k, entries, exits); + loop_manager->get_loop_info(id)->register_handler(0.f); + }; + + if (block_size_k != k) { + mark_k_blocking(); + } else { + brgemm->set_beta(0.f); + } + if (block_size_n != n) + mark_n_blocking(); + if (block_size_m != m) + mark_m_blocking(); modified = true; } diff --git a/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.hpp b/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.hpp index 483a2c5ba53..f51029b0895 100644 --- a/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.hpp +++ b/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.hpp @@ -25,7 +25,11 @@ public: snippets::lowered::LinearIR::constExprIt end) override; private: - static void move_new_memory_buffer(snippets::lowered::LinearIR& linear_ir, const snippets::lowered::LinearIR::constExprIt& brgemm_it); + static snippets::lowered::LinearIR::constExprIt move_new_memory_buffer(snippets::lowered::LinearIR& linear_ir, + const snippets::lowered::LinearIR::constExprIt& brgemm_it); + + static snippets::lowered::LinearIR::constExprIt get_loop_begin_pos(snippets::lowered::LinearIR& linear_ir, + const snippets::lowered::LinearIR::constExprIt& brgemm_it); }; } // namespace pass diff --git a/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/set_brgemm_copy_b_buffers_shape.cpp b/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/set_brgemm_copy_b_buffers_shape.cpp index 701b717fd84..31b135a77da 100644 --- a/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/set_brgemm_copy_b_buffers_shape.cpp +++ b/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/set_brgemm_copy_b_buffers_shape.cpp @@ -28,11 +28,11 @@ bool ov::intel_cpu::pass::SetBrgemmCopyBBuffersShape::run(snippets::lowered::Lin const auto& expr = *expr_it; if (auto copy_b = ov::as_type_ptr(expr->get_node())) { const auto buffer = get_buffer_from_output(expr, 0); - const auto& out_desc = expr->get_output_port_descriptor(0); - buffer->set_allocation_shape(copy_b->get_data_repacking_shape(out_desc->get_shape())); + const auto buffer_shape = copy_b->get_repacking_buffer_shape(); + buffer->set_allocation_shape(buffer_shape); if (copy_b->is_with_compensations()) { const auto compensations_buffer = get_buffer_from_output(expr, 1); - compensations_buffer->set_allocation_shape(copy_b->get_compensation_shape(out_desc->get_shape())); + compensations_buffer->set_allocation_shape(copy_b->get_compensations_buffer_shape()); } modified = true; } diff --git a/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/set_brgemm_cpu_blocking_params.cpp b/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/set_brgemm_cpu_blocking_params.cpp index 41b84b1fe2e..63e33cd4263 100644 --- a/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/set_brgemm_cpu_blocking_params.cpp +++ b/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/set_brgemm_cpu_blocking_params.cpp @@ -50,27 +50,32 @@ pass::SetBrgemmCPUBlockingParams::SetBrgemmCPUBlockingParams() { }; const auto brgemm_in0_dims = snippets::utils::get_planar_pshape(brgemm->input(0)).get_shape(); - const auto M = *(brgemm_in0_dims.rbegin() + 1); - const auto K = *brgemm_in0_dims.rbegin(); const auto brgemm_in1_dims = snippets::utils::get_planar_pshape(brgemm->input(1)).get_shape(); - const auto N = *brgemm_in1_dims.rbegin(); + const auto& M = *++brgemm_in0_dims.rbegin(); + const auto& K = *brgemm_in0_dims.rbegin(); + const auto& N = *brgemm_in1_dims.rbegin(); + const auto m_blk = get_block_size_m(M); + const auto k_blk = get_block_size_k(K); + const auto n_blk = get_block_size_n(N); + if (brgemm->is_with_data_repacking()) { const auto brgemm_copy_b = brgemm->get_brgemm_copy(); - const bool isAMXSupported = dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core_amx); - const auto precision = brgemm_copy_b->get_src_element_type(); const auto brgemmVNNIFactor = brgemm_copy_b->get_brgemm_vnni_factor(); - const bool use_amx = isAMXSupported && precision != ov::element::f32 && (K % brgemmVNNIFactor == 0) && (N % brgemmVNNIFactor == 0); - - const size_t copy_b_block_size_k = use_amx ? get_block_size_k(K) : K; - const size_t copy_b_block_size_n = 64; - - brgemm_copy_b->set_k_block_size(copy_b_block_size_k); - brgemm_copy_b->set_n_block_size(copy_b_block_size_n); + OPENVINO_ASSERT(k_blk == K || k_blk % brgemmVNNIFactor == 0, + "K Block size (", + k_blk, + "), which is not divisible by brgemmVNNIFactor (", + brgemmVNNIFactor, + ") and not equal to K dimension (", + K, + "), is not supported for brgemm data repacking."); + brgemm_copy_b->set_k_block_size(k_blk); + brgemm_copy_b->set_n_block_size(n_blk); } - brgemm->set_m_block_size(get_block_size_m(M)); - brgemm->set_k_block_size(get_block_size_k(K)); - brgemm->set_n_block_size(get_block_size_n(N)); + brgemm->set_m_block_size(m_blk); + brgemm->set_k_block_size(k_blk); + brgemm->set_n_block_size(n_blk); return false; }; diff --git a/src/plugins/intel_cpu/src/transformations/tpp/x64/pass/lowered/set_tpp_leading_dim.cpp b/src/plugins/intel_cpu/src/transformations/tpp/x64/pass/lowered/set_tpp_leading_dim.cpp index 9fabaf4c3f6..7de456d6545 100644 --- a/src/plugins/intel_cpu/src/transformations/tpp/x64/pass/lowered/set_tpp_leading_dim.cpp +++ b/src/plugins/intel_cpu/src/transformations/tpp/x64/pass/lowered/set_tpp_leading_dim.cpp @@ -27,10 +27,10 @@ bool is_planar_layout(const std::vector& layout) { // Directly connected Buffers store data densely, so strides are defined by subternsor dims // Indirectly connected Buffers (with loops between the expr and Buffer) store data according // to their shape and layout -bool has_directly_connected_buffer(const ExpressionPort& port, const snippets::lowered::LinearIR::LoopManagerPtr& loop_mngr) { +bool has_directly_connected_buffer(const ExpressionPort& port, const snippets::lowered::LoopManagerPtr& loop_mngr) { auto accepted_loops = [&loop_mngr, &port](const std::vector& orig, const std::vector& connect) { size_t connect_idx = 0; - auto pred = [&port](const snippets::lowered::LinearIR::LoopManager::LoopPort& loop_port ) { + auto pred = [&port](const snippets::lowered::LoopPort& loop_port ) { return *loop_port.expr_port == port; }; for (const auto orig_loop : orig) { @@ -64,7 +64,7 @@ bool has_directly_connected_buffer(const ExpressionPort& port, const snippets::l return has_buffer; } -size_t get_leading_dim(ExpressionPort port, const snippets::lowered::LinearIR::LoopManagerPtr& loop_mngr) { +size_t get_leading_dim(ExpressionPort port, const snippets::lowered::LoopManagerPtr& loop_mngr) { const auto& port_desc = port.get_descriptor_ptr(); auto layout = port_desc->get_layout(); auto shape = port_desc->get_shape(); diff --git a/src/plugins/intel_cpu/tests/unit/snippets_transformations/lowered/buffer_allocation.cpp b/src/plugins/intel_cpu/tests/unit/snippets_transformations/lowered/buffer_allocation.cpp index b66ce1919f6..764e4ab86ea 100644 --- a/src/plugins/intel_cpu/tests/unit/snippets_transformations/lowered/buffer_allocation.cpp +++ b/src/plugins/intel_cpu/tests/unit/snippets_transformations/lowered/buffer_allocation.cpp @@ -130,6 +130,9 @@ protected: class MHABF16AMXBufferAllocationTest : public BufferAllocationCPUTest { protected: std::shared_ptr GetModel() const override { + const size_t m_blk = 32; + const size_t k_blk = 16; + const size_t n_blk = 64; const auto subtensor_scalar = std::vector{1}; const auto subtensor_power = std::vector{1, ov::snippets::lowered::PortDescriptor::ServiceDimensions::FULL_DIM}; const auto subtensor_full = std::vector(2, ov::snippets::lowered::PortDescriptor::ServiceDimensions::FULL_DIM); @@ -149,9 +152,11 @@ protected: const auto scratch0 = std::make_shared(ov::Shape{ov::intel_cpu::BrgemmCPU::SCRATCH_BYTE_SIZE}); const auto brgemm_cpu0 = std::make_shared( parameter0, brgemm_copyb0->output(0), scratch0, ov::intel_cpu::BrgemmCPU::Type::AMX); - brgemm_cpu0->set_m_block_size(32); - brgemm_cpu0->set_k_block_size(16); - brgemm_cpu0->set_n_block_size(64); + brgemm_cpu0->set_m_block_size(m_blk); + brgemm_cpu0->set_k_block_size(k_blk); + brgemm_copyb0->set_k_block_size(k_blk); + brgemm_cpu0->set_n_block_size(n_blk); + brgemm_copyb0->set_n_block_size(n_blk); const auto relu1 = std::make_shared(brgemm_cpu0); @@ -173,9 +178,11 @@ protected: const auto scratch1 = std::make_shared(ov::Shape{ov::intel_cpu::BrgemmCPU::SCRATCH_BYTE_SIZE}); const auto brgemm_cpu1 = std::make_shared( convert2, brgemm_copyb1->output(0), scratch1, ov::intel_cpu::BrgemmCPU::Type::AMX); - brgemm_cpu1->set_m_block_size(32); - brgemm_cpu1->set_k_block_size(16); - brgemm_cpu1->set_n_block_size(64); + brgemm_cpu1->set_m_block_size(m_blk); + brgemm_cpu1->set_k_block_size(k_blk); + brgemm_copyb1->set_k_block_size(k_blk); + brgemm_cpu1->set_n_block_size(n_blk); + brgemm_copyb1->set_n_block_size(n_blk); const auto relu2 = std::make_shared(brgemm_cpu1); @@ -207,7 +214,7 @@ INSTANTIATE_TEST_SUITE_P(smoke_Snippets_BufferAllocation_MHABF16AMXNotOptimizedW ::testing::Combine( ::testing::Values(false), ::testing::Values(true), - ::testing::Values(196608), + ::testing::Values(167936), ::testing::Values(11)), BufferAllocationCPUTest::getTestCaseName); @@ -215,7 +222,7 @@ INSTANTIATE_TEST_SUITE_P(smoke_Snippets_BufferAllocation_MHAOptimizedWSplit, MHA ::testing::Combine( ::testing::Values(true), ::testing::Values(true), - ::testing::Values(90112), + ::testing::Values(73728), ::testing::Values(3)), BufferAllocationCPUTest::getTestCaseName); @@ -223,7 +230,7 @@ INSTANTIATE_TEST_SUITE_P(smoke_Snippets_BufferAllocation_MHANotOptimizedWOSplit, ::testing::Combine( ::testing::Values(false), ::testing::Values(false), - ::testing::Values(393216), + ::testing::Values(364544), ::testing::Values(11)), BufferAllocationCPUTest::getTestCaseName); @@ -231,8 +238,8 @@ INSTANTIATE_TEST_SUITE_P(smoke_Snippets_BufferAllocation_MHAOptimizedWOSplit, MH ::testing::Combine( ::testing::Values(true), ::testing::Values(false), - ::testing::Values(114688), - ::testing::Values(4)), + ::testing::Values(116736), + ::testing::Values(3)), BufferAllocationCPUTest::getTestCaseName); } // namespace BufferAllocationCPUTest_Instances