natas_hw
/
openvino
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[Snippets] MHA: blocking by K and N for bf16/int8 precisions (#23292) 

### Details:
- *Added brgemm blocking support for bf16 and int8 precisions: in this
case blocking loops are shared between BrgemmCopyB and BrgemmCPU nodes*
- *Reduced allocation shapes of input brgemm buffers in case of low
precision*

### Tickets:
 - *CVS-115165*
This commit is contained in:
Vladislav Golubev 2024-04-25 14:48:00 +02:00 committed by GitHub
parent 1de6329319
commit 6f0e530bce
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
15 changed files with 262 additions and 211 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/common/snippets/src/lowered/pass/cleanup_loop_offsets.cpp
View File

@ -31,13 +31,15 @@ bool CleanupLoopOffsets::run(lowered::LinearIR& linear_ir, lowered::LinearIR::co
}
if (auto outer_loop_end = as_type_ptr<op::LoopEndStatic>(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<PortConnectorPtr, size_t> 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<int64_t>(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.

2
src/common/snippets/src/lowered/pass/fuse_loops.cpp
View File

@ -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;

118
src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.cpp
View File

@ -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<dnnl_data_type_t>(DnnlExtensionUtils::ElementTypeToDataType(m_brgemm_prc_in0));
const auto dt_in1 = static_cast<dnnl_data_type_t>(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<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, 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<matmul::jit_brgemm_matmul_copy_b_t>& 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<int>(N_blk);
brgCopyKernelConf.wei_tag = dnnl_abcd; // What's about other ranks?
brgCopyKernelConf.copy_B_wei_stride = 0;
brgCopyKernelConf.LDB = static_cast<dim_t>(LDB);
brgCopyKernelConf.N = static_cast<dim_t>(N);
brgCopyKernelConf.N_tail = static_cast<dim_t>(N_tail);
brgCopyKernelConf.N_tail = static_cast<dim_t>(N_tail);
brgCopyKernelConf.N_blk = static_cast<dim_t>(N_blk);
brgCopyKernelConf.K = static_cast<dim_t>(K);
brgCopyKernelConf.K_blk = static_cast<dim_t>(K);
@ -91,8 +104,8 @@ void jit_brgemm_copy_b_emitter::init_brgemm_copy(std::unique_ptr<matmul::jit_brg
brgCopyKernelConf.isa = avx512_core_amx;
brgCopyKernelConf.s8s8_compensation_required = false;
} else {
brgCopyKernelConf.isa = dt_in0 == dnnl_data_type_t::dnnl_bf16 ? avx512_core_bf16 : avx512_core_vnni;
brgCopyKernelConf.s8s8_compensation_required = dt_in0 == dnnl_data_type_t::dnnl_s8;
brgCopyKernelConf.isa = src_dt == dnnl_data_type_t::dnnl_bf16 ? avx512_core_bf16 : avx512_core_vnni;
brgCopyKernelConf.s8s8_compensation_required = src_dt == dnnl_data_type_t::dnnl_s8;
}
brgCopyKernelConf.has_zero_point_a = false;
@ -100,41 +113,38 @@ void jit_brgemm_copy_b_emitter::init_brgemm_copy(std::unique_ptr<matmul::jit_brg
brgCopyKernelConf.src_zp_type = dnnl::impl::cpu::x64::none;
auto status = matmul::create_brgemm_matmul_copy_b(kernel, &brgCopyKernelConf);
if (status != dnnl_success)
OV_CPU_JIT_EMITTER_THROW("cannot create kernel due to invalid params");
OV_CPU_JIT_EMITTER_ASSERT(status == dnnl_success, "cannot create kernel due to invalid params");
}
void jit_brgemm_copy_b_emitter::emit_impl(const std::vector<size_t>& in,
const std::vector<size_t>& out) const {
if (host_isa_ == cpu::x64::avx512_core) {
Xbyak::Reg64 src(static_cast<int>(in[0]));
Xbyak::Reg64 dst(static_cast<int>(out[0]));
Xbyak::Reg64 comp(static_cast<int>(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<int>(out[1]));
}
void jit_brgemm_copy_b_emitter::validate_arguments(const std::vector<size_t> &in, const std::vector<size_t> &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<size_t>& in, const std::vector<size_t>& 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<int>(in[0]));
Xbyak::Reg64 dst(static_cast<int>(out[0]));
Xbyak::Reg64 comp(static_cast<int>(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);
}

19
src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_copy_b_emitter.hpp
View File

@ -23,6 +23,7 @@ public:
}
private:
void validate_arguments(const std::vector<size_t> &in, const std::vector<size_t> &out) const override;
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,
@ -36,18 +37,24 @@ private:
const void* src, const void* dst, const void* comp, size_t N, size_t K);
std::unique_ptr<dnnl::impl::cpu::x64::matmul::jit_brgemm_matmul_copy_b_t> 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

10
src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_emitter.cpp
View File

@ -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<dnnl_data_type_t>(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<brge
status = brgemm_init_tiles(desc, ctx.palette);
ctx.is_with_comp = ctx.dt_in0 == dnnl_data_type_t::dnnl_s8 && !ctx.is_with_amx;
brgemm_kernel_t* kernel_ = nullptr;
status = brgemm_kernel_create(&kernel_, desc);
if (status != dnnl_success)

1
src/plugins/intel_cpu/src/emitters/snippets/x64/jit_brgemm_emitter.hpp
View File

@ -58,7 +58,6 @@ private:
std::unique_ptr<dnnl::impl::cpu::x64::brgemm_kernel_t> 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;

19
src/plugins/intel_cpu/src/emitters/snippets/x64/verbose.cpp
View File

@ -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();
}

23
src/plugins/intel_cpu/src/transformations/snippets/x64/op/brgemm_copy_b.cpp
View File

@ -18,7 +18,6 @@ intel_cpu::BrgemmCopyB::BrgemmCopyB(const Output<Node>& x, const element::Type s
std::vector<size_t> 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<Node>& 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<Node>& x, const element::Type src_type, const Type type,
@ -34,7 +35,6 @@ intel_cpu::BrgemmCopyB::BrgemmCopyB(const Output<Node>& x, const element::Type s
std::vector<size_t> 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<Node>& 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<Node> intel_cpu::BrgemmCopyB::clone_with_new_inputs(const OutputVector& new_args) const {

8
src/plugins/intel_cpu/src/transformations/snippets/x64/op/brgemm_copy_b.hpp
View File

@ -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;
};

185
src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.cpp
View File

@ -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<ov::snippets::op::NewMemoryBuffer>(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<snippets::op::Brgemm>(node);
const auto brgemm_cpu = ov::as_type_ptr<intel_cpu::BrgemmCPU>(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<LoopPort> 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<LoopPort> 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<LoopPort> 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<LoopPort> 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<LoopPort> 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<LoopPort> 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<ov::snippets::lowered::SpecificIterationHandlers::HandlerType::FIRST_ITER, SetBrgemmBeta>(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<LoopPort> 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<LoopPort> 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<LoopPort> 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<LoopPort> 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<LoopPort> 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<LoopPort> 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<SpecificIterationHandlers::HandlerType::FIRST_ITER, SetBrgemmBeta>(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;
}

6
src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.hpp
View File

@ -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

6
src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/set_brgemm_copy_b_buffers_shape.cpp
View File

@ -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<ov::intel_cpu::BrgemmCopyB>(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;
}

35
src/plugins/intel_cpu/src/transformations/snippets/x64/pass/set_brgemm_cpu_blocking_params.cpp
View File

@ -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;
};

6
src/plugins/intel_cpu/src/transformations/tpp/x64/pass/lowered/set_tpp_leading_dim.cpp
View File

@ -27,10 +27,10 @@ bool is_planar_layout(const std::vector<size_t>& 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<size_t>& orig, const std::vector<size_t>& 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();

29
src/plugins/intel_cpu/tests/unit/snippets_transformations/lowered/buffer_allocation.cpp
View File

@ -130,6 +130,9 @@ protected:
class MHABF16AMXBufferAllocationTest : public BufferAllocationCPUTest {
protected:
std::shared_ptr<ov::Model> 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<size_t>{1};
const auto subtensor_power = std::vector<size_t>{1, ov::snippets::lowered::PortDescriptor::ServiceDimensions::FULL_DIM};
const auto subtensor_full = std::vector<size_t>(2, ov::snippets::lowered::PortDescriptor::ServiceDimensions::FULL_DIM);
@ -149,9 +152,11 @@ protected:
const auto scratch0 = std::make_shared<ov::snippets::op::NewMemoryBuffer>(ov::Shape{ov::intel_cpu::BrgemmCPU::SCRATCH_BYTE_SIZE});
const auto brgemm_cpu0 = std::make_shared<ov::intel_cpu::BrgemmCPU>(
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<ov::op::v0::Relu>(brgemm_cpu0);
@ -173,9 +178,11 @@ protected:
const auto scratch1 = std::make_shared<ov::snippets::op::NewMemoryBuffer>(ov::Shape{ov::intel_cpu::BrgemmCPU::SCRATCH_BYTE_SIZE});
const auto brgemm_cpu1 = std::make_shared<ov::intel_cpu::BrgemmCPU>(
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<ov::op::v0::Relu>(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