[Snippets][AMX] Perform amx_tile_configure only when necessary (#23016)
### Details: - Currently, `amx_tile_configure` in brgemm emitter is called for each inference whereas reconfiguration can be called only when input shapes are differ from previous inference. This PR introduces `amx_tile_config` structure with the latest used dimensions. The structure is passed as a runtime argument in compiled kernel. Brgemm emitter gets the structure and checks if the current shapes are the same as previous. If they are the same, `amx_tile_configure` is skipped. Otherwise, the reconfiguration is called, and `amx_tile_config` is updated with new shapes. ### Tickets: - *CVS-125676*
This commit is contained in:
Vladislav Golubev
GitHub
parent
5bc33d8970
commit
dfc39c0d17
|
|
@ -606,11 +606,12 @@ TokenizeSnippets::TokenizeSnippets() {
|
|||
}
|
||||
|
||||
// The each data node (Parameter (and non-Scalar Constants), Result, Buffers with the same ID) requires the own unique GPR.
|
||||
// At the moment, CPU Plugin has limitation for GPR registers: there are only 12 available registers.
|
||||
// At the moment, CPU Plugin has limitation for GPR registers: there are 12 available GPRs,
|
||||
// and one of them must be reserved for runtime parameters, so only 11 can be used during kernel execution.
|
||||
// This limitation will be resolved once generator supports gprs spills [75622].
|
||||
// TODO [75567]: move this plugin-specific constraint to the plugin callback
|
||||
const auto unique_buffer_count = op::Subgraph::get_estimated_buffer_count(ops_for_buffer_count);
|
||||
if (body_parameters.size() + body_results.size() + hidden_data_count + unique_buffer_count > 12) {
|
||||
if (body_parameters.size() + body_results.size() + hidden_data_count + unique_buffer_count > 11) {
|
||||
const std::string message_reset = "new subgraph is created. Impossible to schedule subgraph with " +
|
||||
std::to_string(body_parameters.size()) + " inputs, " + std::to_string(body_results.size()) + " outputs and " +
|
||||
std::to_string(hidden_data_count) + " non-scalar constants and " + std::to_string(unique_buffer_count) + "buffers.";
|
||||
|
|
|
|||
|
|
@ -454,7 +454,7 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken
|
|||
|
||||
// TODO [75567]: move this plugin-specific constraint to the plugin callback
|
||||
const auto last_node = ordered_ops.back();
|
||||
if (potential_body_params_count + last_node->get_output_size() + hidden_virtual_ports_count + buffer_count > 12) {
|
||||
if (potential_body_params_count + last_node->get_output_size() + hidden_virtual_ports_count + buffer_count > 11) {
|
||||
return false;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -13,7 +13,6 @@
|
|||
|
||||
#include <cpu/x64/brgemm/brgemm.hpp>
|
||||
#include <cpu/x64/matmul/brgemm_matmul_utils.hpp>
|
||||
#include <cpu/x64/amx_tile_configure.hpp>
|
||||
|
||||
|
||||
using namespace Xbyak;
|
||||
|
|
@ -136,34 +135,6 @@ void jit_brgemm_copy_b_emitter::emit_impl(const std::vector<size_t>& in,
|
|||
|
||||
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 {
|
||||
constexpr size_t gpr_size = 8;
|
||||
Xbyak::Operand gprs_to_save[] = {h->r8, h->r9, h->r10, h->r11, h->r12, h->r13, h->r14, h->r15,
|
||||
h->rax, h->rcx, h->rdx, h->rdi, h->rsi, h->rbp, h->rbx};
|
||||
size_t n_gprs_to_save = sizeof(gprs_to_save) / sizeof(gprs_to_save[0]);
|
||||
|
||||
h->sub(h->rsp, n_gprs_to_save * gpr_size);
|
||||
for (size_t i = 0; i < n_gprs_to_save; ++i)
|
||||
h->mov(h->ptr[h->rsp + i * gpr_size], gprs_to_save[i]);
|
||||
|
||||
// caller obligation to save k-regs as callee may use them
|
||||
size_t n_k_regs_to_save = 8;
|
||||
h->sub(h->rsp, n_k_regs_to_save * k_mask_size);
|
||||
for (size_t i = 0; i < n_k_regs_to_save; ++i) {
|
||||
if (mayiuse(avx512_core))
|
||||
h->kmovq(h->ptr[h->rsp + i * k_mask_size], Opmask(static_cast<int>(i)));
|
||||
else
|
||||
h->kmovw(h->ptr[h->rsp + i * k_mask_size], Opmask(static_cast<int>(i)));
|
||||
}
|
||||
|
||||
// 1. Caller obligation to save vector registers as callee may use them.
|
||||
// 2. There is an implicit assumption that the host code uses the same
|
||||
// `isa` as the injector. Once the assumption is wrong, `vecs_count` and
|
||||
// `vlen` should be replaced with `host_isa::vlen` and
|
||||
// `host_isa::vecs_count`.
|
||||
h->sub(h->rsp, get_max_vecs_count() * get_vec_length());
|
||||
for (size_t i = 0; i < get_max_vecs_count(); ++i)
|
||||
h->uni_vmovups(h->ptr[h->rsp + i * get_vec_length()], Zmm(i));
|
||||
|
||||
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);
|
||||
|
|
@ -176,6 +147,7 @@ void jit_brgemm_copy_b_emitter::emit_kernel_call(const matmul::jit_brgemm_matmul
|
|||
};
|
||||
#endif
|
||||
|
||||
internal_call_preamble();
|
||||
// save function address in gpr to pass in call instruction
|
||||
const auto &kernel_overload = static_cast<void (*)(matmul::jit_brgemm_matmul_copy_b_t*,
|
||||
const void*,
|
||||
|
|
@ -216,38 +188,15 @@ void jit_brgemm_copy_b_emitter::emit_kernel_call(const matmul::jit_brgemm_matmul
|
|||
h->mov(abi_param5, N);
|
||||
h->mov(abi_param6, K);
|
||||
#endif
|
||||
// align stack on 16-byte as ABI requires
|
||||
// note that RBX must not be changed by the callee
|
||||
h->mov(h->rbx, h->rsp);
|
||||
h->and_(h->rbx, 0xf);
|
||||
h->sub(h->rsp, h->rbx);
|
||||
|
||||
internal_call_rsp_align();
|
||||
h->call(h->rbp);
|
||||
|
||||
h->add(h->rsp, h->rbx);
|
||||
internal_call_rsp_restore();
|
||||
|
||||
#ifdef _WIN32
|
||||
h->add(h->rsp, gpr_size * num_args_passed_on_stack);
|
||||
#endif
|
||||
// restore vector registers
|
||||
for (int i = static_cast<int>(get_max_vecs_count()) - 1; i >= 0; --i) {
|
||||
h->uni_vmovups(Zmm(i), h->ptr[h->rsp + i * get_vec_length()]);
|
||||
}
|
||||
h->add(h->rsp, (get_max_vecs_count()) * get_vec_length());
|
||||
|
||||
// restore k registers
|
||||
for (int i = n_k_regs_to_save - 1; i >= 0; --i) {
|
||||
if (mayiuse(avx512_core))
|
||||
h->kmovq(Opmask(i), h->ptr[h->rsp + i * k_mask_size]);
|
||||
else
|
||||
h->kmovw(Opmask(i), h->ptr[h->rsp + i * k_mask_size]);
|
||||
}
|
||||
h->add(h->rsp, n_k_regs_to_save * k_mask_size);
|
||||
|
||||
// restore gpr registers
|
||||
for (int i = n_gprs_to_save - 1; i >= 0; --i)
|
||||
h->mov(gprs_to_save[i], h->ptr[h->rsp + i * gpr_size]);
|
||||
h->add(h->rsp, n_gprs_to_save * gpr_size);
|
||||
internal_call_postamble();
|
||||
}
|
||||
|
||||
void jit_brgemm_copy_b_emitter::execute(matmul::jit_brgemm_matmul_copy_b_t *kernel, const void *src,
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
//
|
||||
|
||||
#include "jit_brgemm_emitter.hpp"
|
||||
#include "jit_kernel_emitter.hpp"
|
||||
|
||||
#include "snippets/utils.hpp"
|
||||
#include "snippets/lowered/expression.hpp"
|
||||
|
|
@ -87,12 +88,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);
|
||||
bool with_amx = brgemm_node->is_amx();
|
||||
|
||||
io_data_size = {brg0Prc.size(), brg1Prc.size()};
|
||||
if (brgemm_node->get_input_size() == 3)
|
||||
io_data_size.push_back(brgemm_node->get_input_element_type(2).size());
|
||||
io_data_size.push_back(brgemm_node->get_output_element_type(0).size());
|
||||
|
||||
m_with_comp = brgemm_node->is_with_compensations();
|
||||
m_with_scratch = brgemm_node->is_with_scratchpad();
|
||||
|
|
@ -117,8 +112,9 @@ jit_brgemm_emitter::jit_brgemm_emitter(jit_generator* h, cpu_isa_t isa, const ov
|
|||
m_ctx.dt_in0 = static_cast<dnnl_data_type_t>(DnnlExtensionUtils::ElementTypeToDataType(brg0Prc));
|
||||
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();
|
||||
|
||||
init_brgemm_kernel(m_ctx, m_kernel, with_amx);
|
||||
init_brgemm_kernel(m_ctx, m_kernel);
|
||||
|
||||
m_load_offset_a = brgemm_node->get_offset_a();
|
||||
m_load_offset_b = brgemm_node->get_offset_b();
|
||||
|
|
@ -147,20 +143,17 @@ std::set<std::vector<element::Type>> jit_brgemm_emitter::get_supported_precision
|
|||
}
|
||||
}
|
||||
|
||||
void jit_brgemm_emitter::init_brgemm_kernel(brgemmCtx& ctx, std::unique_ptr<brgemm_kernel_t>& kernel, bool use_amx) {
|
||||
void jit_brgemm_emitter::init_brgemm_kernel(brgemmCtx& ctx, std::unique_ptr<brgemm_kernel_t>& kernel) {
|
||||
brgemm_t desc;
|
||||
const bool is_int8 = utils::one_of(ctx.dt_in0, data_type::u8, data_type::s8) && utils::one_of(ctx.dt_in1, data_type::u8, data_type::s8);
|
||||
auto isa = use_amx ? isa_undef
|
||||
: ctx.dt_in0 == dnnl_data_type_t::dnnl_bf16 ? avx512_core_bf16 : (is_int8 ? avx512_core_vnni : avx512_core);
|
||||
auto isa = ctx.is_with_amx ? avx512_core_amx
|
||||
: ctx.dt_in0 == dnnl_data_type_t::dnnl_bf16 ? avx512_core_bf16 : (is_int8 ? avx512_core_vnni : avx512_core);
|
||||
auto status = brgemm_desc_init(&desc, isa, brgemm_strd, ctx.dt_in0, ctx.dt_in1,
|
||||
false, false, brgemm_row_major, 1.f, ctx.beta, ctx.LDA, ctx.LDB, ctx.LDC, ctx.M, ctx.N, ctx.K, nullptr);
|
||||
if (status != dnnl_success)
|
||||
OV_CPU_JIT_EMITTER_THROW("cannot initialize brgemm descriptor due to invalid params");
|
||||
|
||||
ctx.is_with_amx = use_amx;
|
||||
status = brgemm_init_tiles(desc, ctx.palette);
|
||||
if (use_amx)
|
||||
amx_tile_configure(ctx.palette);
|
||||
|
||||
ctx.is_with_comp = ctx.dt_in0 == dnnl_data_type_t::dnnl_s8 && !ctx.is_with_amx;
|
||||
|
||||
|
|
@ -184,7 +177,6 @@ void jit_brgemm_emitter::emit_impl(const std::vector<size_t>& in, const std::vec
|
|||
Xbyak::Reg64 input_2(static_cast<int>(m_with_scratch ? in[2] : 0)); // scratch. Default reg index is 0 if there isn't scratch
|
||||
Xbyak::Reg64 output_0(static_cast<int>(out[0]));
|
||||
emit_brgemm_kernel_call(m_kernel.get(),
|
||||
m_ctx,
|
||||
input_0,
|
||||
input_1,
|
||||
input_2,
|
||||
|
|
@ -198,117 +190,80 @@ void jit_brgemm_emitter::emit_impl(const std::vector<size_t>& in, const std::vec
|
|||
}
|
||||
}
|
||||
|
||||
void jit_brgemm_emitter::emit_brgemm_kernel_call(const brgemm_kernel_t *brg_kernel, const brgemmCtx& ctx,
|
||||
void jit_brgemm_emitter::emit_brgemm_kernel_call(const brgemm_kernel_t *brg_kernel,
|
||||
Reg64 addr_A, Reg64 addr_B, Reg64 scratch, Reg64 addr_C,
|
||||
const size_t in0_kernel_offset, const size_t in1_kernel_offset,
|
||||
const size_t in2_kernel_offset, const size_t out0_kernel_offset) const {
|
||||
constexpr size_t gpr_size = 8;
|
||||
if (ctx.is_with_amx) {
|
||||
Xbyak::Operand gprs_to_save[] = {h->r8, h->r9, h->r10, h->r11, h->rax,
|
||||
h->rcx, h->rdx, h->rdi, h->rsi, h->rbp, h->rbx};
|
||||
size_t n_gprs_to_save = sizeof(gprs_to_save) / sizeof(gprs_to_save[0]);
|
||||
|
||||
h->sub(h->rsp, n_gprs_to_save * gpr_size);
|
||||
for (size_t i = 0; i < n_gprs_to_save; ++i)
|
||||
h->mov(h->ptr[h->rsp + i * gpr_size], gprs_to_save[i]);
|
||||
|
||||
// save function address in gpr to pass in call instruction
|
||||
const auto& overload = static_cast<status_t(*)(const char*)>(amx_tile_configure);
|
||||
h->mov(h->rbp, reinterpret_cast<uintptr_t>(overload));
|
||||
h->mov(abi_param1, reinterpret_cast<uintptr_t>(ctx.palette));
|
||||
|
||||
// align stack on 16-byte as ABI requires
|
||||
// note that RBX must not be changed by the callee
|
||||
h->mov(h->rbx, h->rsp);
|
||||
h->and_(h->rbx, 0xf);
|
||||
h->sub(h->rsp, h->rbx);
|
||||
|
||||
h->call(h->rbp);
|
||||
|
||||
h->add(h->rsp, h->rbx);
|
||||
// restore gpr registers
|
||||
for (int i = n_gprs_to_save - 1; i >= 0; --i)
|
||||
h->mov(gprs_to_save[i], h->ptr[h->rsp + i * gpr_size]);
|
||||
h->add(h->rsp, n_gprs_to_save * gpr_size);
|
||||
}
|
||||
|
||||
internal_call_preamble();
|
||||
|
||||
// save function address in gpr to pass in call instruction
|
||||
const auto& brgemm_kernel_overload = static_cast<void (*)(const brgemm_kernel_t*,
|
||||
const void*,
|
||||
const void*,
|
||||
void*,
|
||||
void*,
|
||||
int)>(kernel_execute);
|
||||
const brgemmCtx*,
|
||||
jit_brgemm_call_args*)>(kernel_execute);
|
||||
h->mov(h->rbp, reinterpret_cast<uintptr_t>(brgemm_kernel_overload));
|
||||
// todo: several of addr_{A, B, C} could be also abi_paramX, so one of them could be corrupted
|
||||
// if moving directly h->uni_vmovq(abi_paramX, adr_X). Save them to vector regs to avoid corruption.
|
||||
// It's likely that a more efficient solution exists.
|
||||
h->uni_vmovq(Xmm(0), addr_A);
|
||||
h->uni_vmovq(Xmm(1), addr_B);
|
||||
h->uni_vmovq(Xmm(2), addr_C);
|
||||
if (m_with_scratch)
|
||||
h->uni_vmovq(Xmm(3), scratch);
|
||||
// todo: Windows ABI : requires different num of arguments passed in regs and on the stack. Need to align.
|
||||
const auto data_ptr_reg = [&](Xmm xmm, Xbyak::Reg64 reg, size_t bytes_offset) {
|
||||
h->uni_vmovq(reg, xmm);
|
||||
if (bytes_offset) h->add(reg, bytes_offset);
|
||||
};
|
||||
h->mov(abi_param1, reinterpret_cast<uintptr_t>(brg_kernel));
|
||||
data_ptr_reg(Xmm(0), abi_param2, in0_kernel_offset);
|
||||
data_ptr_reg(Xmm(1), abi_param3, in1_kernel_offset);
|
||||
data_ptr_reg(Xmm(2), abi_param4, out0_kernel_offset);
|
||||
|
||||
auto reserved_stack_size = sizeof(jit_brgemm_call_args);
|
||||
#ifdef _WIN32
|
||||
// Before function call we should allocate stack area for
|
||||
// - register parameters - ABI parameters (shadow space)
|
||||
// - stack parameters - remaining parameters
|
||||
const size_t num_args_passed_on_stack = 6; // count of function brgemm_kernel_overload() parameters
|
||||
size_t abi_param_count = sizeof(abi_param_regs) / sizeof(abi_param_regs[0]);
|
||||
h->sub(h->rsp, num_args_passed_on_stack * gpr_size);
|
||||
|
||||
// Push the remaining parameters on the stack
|
||||
if (m_with_scratch) {
|
||||
h->uni_vmovq(h->qword[h->rsp + (abi_param_count + 0) * gpr_size], Xmm(3));
|
||||
if (in2_kernel_offset) h->add(h->qword[h->rsp + (abi_param_count + 0) * gpr_size], in2_kernel_offset);
|
||||
} else {
|
||||
h->mov(h->qword[h->rsp + (abi_param_count + 0) * gpr_size], reinterpret_cast<uintptr_t>(nullptr));
|
||||
}
|
||||
h->mov(abi_not_param1, static_cast<int>(m_with_comp));
|
||||
h->mov(h->qword[h->rsp + (abi_param_count + 1) * gpr_size], abi_not_param1);
|
||||
#else
|
||||
if (m_with_scratch) {
|
||||
data_ptr_reg(Xmm(3), abi_param5, in2_kernel_offset);
|
||||
} else {
|
||||
h->mov(abi_param5, reinterpret_cast<uintptr_t>(nullptr));
|
||||
}
|
||||
h->mov(abi_param6, static_cast<int>(m_with_comp));
|
||||
// Before function call we should also allocate stack area for ABI parameters (shadow space)
|
||||
reserved_stack_size += 3 * gpr_size;
|
||||
#endif
|
||||
|
||||
// Reserve memory on the stack
|
||||
h->sub(h->rsp, reserved_stack_size);
|
||||
|
||||
auto write_addr_on_stack = [&](size_t arg_offset, Reg64 addr, size_t addr_offset) {
|
||||
const auto stack_frame = h->qword[h->rsp + arg_offset];
|
||||
h->mov(stack_frame, addr);
|
||||
if (addr_offset) h->add(stack_frame, addr_offset);
|
||||
};
|
||||
|
||||
write_addr_on_stack(GET_OFF_BRGEMM_ARGS(A), addr_A, in0_kernel_offset);
|
||||
write_addr_on_stack(GET_OFF_BRGEMM_ARGS(B), addr_B, in1_kernel_offset);
|
||||
write_addr_on_stack(GET_OFF_BRGEMM_ARGS(C), addr_C, out0_kernel_offset);
|
||||
|
||||
if (m_with_scratch) {
|
||||
write_addr_on_stack(GET_OFF_BRGEMM_ARGS(scratch), scratch, in2_kernel_offset);
|
||||
} else {
|
||||
h->mov(h->qword[h->rsp + GET_OFF_BRGEMM_ARGS(scratch)], reinterpret_cast<uintptr_t>(nullptr));
|
||||
}
|
||||
|
||||
// abi_param1 always contains jit_snippets_call_args which has amx tile config for each thread
|
||||
h->lea(h->r10, h->ptr[abi_param1 + GET_OFF(amx_tile_config)]);
|
||||
h->mov(h->qword[h->rsp + GET_OFF_BRGEMM_ARGS(amx_tile_config)], h->r10);
|
||||
|
||||
h->mov(abi_param1, reinterpret_cast<uintptr_t>(brg_kernel));
|
||||
h->mov(abi_param2, reinterpret_cast<uintptr_t>(&m_ctx));
|
||||
h->mov(abi_param3, h->rsp);
|
||||
|
||||
internal_call_rsp_align();
|
||||
h->call(h->rbp);
|
||||
internal_call_rsp_restore();
|
||||
|
||||
#ifdef _WIN32
|
||||
h->add(h->rsp, num_args_passed_on_stack * gpr_size);
|
||||
#endif
|
||||
|
||||
h->add(h->rsp, reserved_stack_size);
|
||||
internal_call_postamble();
|
||||
}
|
||||
|
||||
void jit_brgemm_emitter::kernel_execute(const brgemm_kernel_t *brg_kernel, const void *A, const void *B, void *C, void *scratch, int with_comp) {
|
||||
void jit_brgemm_emitter::kernel_execute(const dnnl::impl::cpu::x64::brgemm_kernel_t* brg_kernel,
|
||||
const brgemmCtx* ctx,
|
||||
jit_brgemm_call_args* brgemm_call_args) {
|
||||
if (ctx->is_with_amx) {
|
||||
const auto& amx_tile_config = brgemm_call_args->amx_tile_config;
|
||||
if (ctx->M != amx_tile_config->M || ctx->K != amx_tile_config->K || ctx->N != amx_tile_config->N) {
|
||||
amx_tile_config->M = ctx->M;
|
||||
amx_tile_config->K = ctx->K;
|
||||
amx_tile_config->N = ctx->N;
|
||||
amx_tile_configure(ctx->palette);
|
||||
}
|
||||
}
|
||||
|
||||
brgemm_kernel_params_t brgemm_p;
|
||||
|
||||
brgemm_p.batch = nullptr; // default value
|
||||
brgemm_p.ptr_A = A;
|
||||
brgemm_p.ptr_B = B;
|
||||
brgemm_p.ptr_C = C;
|
||||
brgemm_p.ptr_D = C;
|
||||
brgemm_p.ptr_buf = scratch;
|
||||
brgemm_p.ptr_A = brgemm_call_args->A;
|
||||
brgemm_p.ptr_B = brgemm_call_args->B;
|
||||
brgemm_p.ptr_C = brgemm_call_args->C;
|
||||
brgemm_p.ptr_D = brgemm_call_args->C;
|
||||
brgemm_p.ptr_buf = brgemm_call_args->scratch;
|
||||
brgemm_p.ptr_bias = nullptr;
|
||||
brgemm_p.do_post_ops = static_cast<size_t>(with_comp);
|
||||
brgemm_p.do_apply_comp = static_cast<size_t>(with_comp);
|
||||
brgemm_p.do_post_ops = static_cast<size_t>(ctx->is_with_comp);
|
||||
brgemm_p.do_apply_comp = static_cast<size_t>(ctx->is_with_comp);
|
||||
brgemm_p.skip_accm = 0;
|
||||
brgemm_p.BS = 1; // default value
|
||||
OV_CPU_JIT_EMITTER_ASSERT(brg_kernel != nullptr, "has nullptr kernel");
|
||||
|
|
|
|||
|
|
@ -5,13 +5,22 @@
|
|||
#pragma once
|
||||
|
||||
#include "emitters/plugin/x64/jit_emitter.hpp"
|
||||
#include "jit_kernel_emitter.hpp"
|
||||
|
||||
#include <cpu/x64/brgemm/brgemm.hpp>
|
||||
|
||||
|
||||
namespace ov {
|
||||
namespace intel_cpu {
|
||||
|
||||
struct jit_brgemm_call_args {
|
||||
const void* A = nullptr;
|
||||
const void* B = nullptr;
|
||||
void* C = nullptr;
|
||||
void* scratch = nullptr;
|
||||
amx_tile_config_t* amx_tile_config = nullptr;
|
||||
};
|
||||
#define GET_OFF_BRGEMM_ARGS(field) offsetof(jit_brgemm_call_args, field)
|
||||
|
||||
class jit_brgemm_emitter : public jit_emitter {
|
||||
public:
|
||||
jit_brgemm_emitter(dnnl::impl::cpu::x64::jit_generator* h, dnnl::impl::cpu::x64::cpu_isa_t isa, const ov::snippets::lowered::ExpressionPtr& expr);
|
||||
|
|
@ -35,13 +44,15 @@ private:
|
|||
bool is_with_comp;
|
||||
float beta;
|
||||
};
|
||||
static void init_brgemm_kernel(brgemmCtx& ctx, std::unique_ptr<dnnl::impl::cpu::x64::brgemm_kernel_t>& brgKernel, bool use_amx);
|
||||
static void init_brgemm_kernel(brgemmCtx& ctx, std::unique_ptr<dnnl::impl::cpu::x64::brgemm_kernel_t>& brgKernel);
|
||||
|
||||
void emit_brgemm_kernel_call(const dnnl::impl::cpu::x64::brgemm_kernel_t* brg_kernel, const brgemmCtx& ctx,
|
||||
void emit_brgemm_kernel_call(const dnnl::impl::cpu::x64::brgemm_kernel_t* brg_kernel,
|
||||
Xbyak::Reg64 addr_A, Xbyak::Reg64 addr_B, Xbyak::Reg64 scratch, Xbyak::Reg64 addr_C,
|
||||
size_t in0_kernel_offset = 0, size_t in1_kernel_offset = 0,
|
||||
size_t in2_kernel_offset = 0, size_t out0_kernel_offset = 0) const;
|
||||
static void kernel_execute(const dnnl::impl::cpu::x64::brgemm_kernel_t *brg_kernel, const void *A, const void *B, void *C, void *scratch, int with_comp);
|
||||
static void kernel_execute(const dnnl::impl::cpu::x64::brgemm_kernel_t* brg_kernel,
|
||||
const brgemmCtx* ctx,
|
||||
jit_brgemm_call_args* brgemm_call_args);
|
||||
|
||||
brgemmCtx m_ctx;
|
||||
std::unique_ptr<dnnl::impl::cpu::x64::brgemm_kernel_t> m_kernel = nullptr;
|
||||
|
|
@ -54,8 +65,6 @@ private:
|
|||
size_t m_load_offset_scratch = 0lu;
|
||||
size_t m_store_offset_c = 0lu;
|
||||
|
||||
std::vector<size_t> io_data_size {};
|
||||
|
||||
#ifdef SNIPPETS_DEBUG_CAPS
|
||||
friend std::string init_info_jit_brgemm_emitter(const jit_brgemm_emitter *emitter);
|
||||
#endif
|
||||
|
|
|
|||
|
|
@ -163,8 +163,8 @@ void jit_kernel_emitter::init_body_regs(const std::set<size_t>& kernel_regs,
|
|||
for (const auto& abstract_to_physical : gpr_map_pool.first)
|
||||
data_ptr_regs_idx.push_back(abstract_to_physical.second);
|
||||
|
||||
vec_map_pool.second.insert(vec_map_pool.second.end(), pool_vec_idxs.cbegin(), pool_vec_idxs.cend());
|
||||
gpr_map_pool.second.insert(gpr_map_pool.second.end(), pool_gpr_idxs.cbegin(), pool_gpr_idxs.cend());
|
||||
vec_map_pool.second.insert(vec_map_pool.second.end(), pool_vec_idxs.cbegin(), pool_vec_idxs.cend());
|
||||
map_abstract_registers(gpr_map_pool, vec_map_pool, general_exprs);
|
||||
}
|
||||
|
||||
|
|
@ -231,9 +231,9 @@ jit_kernel_static_emitter::jit_kernel_static_emitter(dnnl::impl::cpu::x64::jit_g
|
|||
master_shape.insert(master_shape.begin(), jcp.parallel_executor_ndims - master_shape.size(), 1);
|
||||
|
||||
// - Reserve abi_param1 and abi_param2, since they'll be used to pass runtime call args to kernel
|
||||
// - However we can use reg_indexes_idx and reg_runtime_params_idx for non memory access operations
|
||||
// - However we can use reg_indexes_idx for non memory access operations
|
||||
// since we won't need them after offsets calculation
|
||||
init_body_regs({reg_indexes_idx, reg_runtime_params_idx}, {}, {reg_indexes_idx, reg_runtime_params_idx});
|
||||
init_body_regs({reg_indexes_idx, reg_runtime_params_idx}, {}, {reg_indexes_idx});
|
||||
}
|
||||
|
||||
void jit_kernel_static_emitter::init_data_pointers(const std::vector<Xbyak::Reg64>& data_ptr_regs) const {
|
||||
|
|
|
|||
|
|
@ -20,6 +20,12 @@ namespace intel_cpu {
|
|||
#define GET_OFF(field) offsetof(jit_snippets_call_args, field)
|
||||
#define GET_OFF_LOOP_ARGS(field) offsetof(jit_snippets_call_args::loop_args_t, field)
|
||||
|
||||
struct amx_tile_config_t {
|
||||
size_t M = 0;
|
||||
size_t K = 0;
|
||||
size_t N = 0;
|
||||
};
|
||||
|
||||
struct jit_snippets_call_args {
|
||||
struct loop_args_t;
|
||||
|
||||
|
|
@ -31,12 +37,12 @@ struct jit_snippets_call_args {
|
|||
const void *src_ptrs[SNIPPETS_MAX_SNIPPETS_DIMS] = {};
|
||||
void *dst_ptrs[SNIPPETS_MAX_SNIPPETS_DIMS] = {};
|
||||
void *buffer_scratchpad_ptr = nullptr;
|
||||
|
||||
// Note: Ideally loop_args must be private, since we manage this pointer manually.
|
||||
// However, standard-layout class definition (to use offset_of) requires the same access specifier
|
||||
// for all non-static data members. So we can keep them public or friend all control-flow emitters
|
||||
int32_t num_loops = 0;
|
||||
loop_args_t* loop_args = nullptr;
|
||||
amx_tile_config_t amx_tile_config;
|
||||
};
|
||||
|
||||
struct jit_snippets_call_args::loop_args_t {
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ std::string MaxNumParamsEltwise::getTestCaseName(testing::TestParamInfo<ov::test
|
|||
void MaxNumParamsEltwise::SetUp() {
|
||||
ov::test::InputShape inputShape;
|
||||
std::tie(inputShape, ref_num_nodes, ref_num_subgraphs, targetDevice) = this->GetParam();
|
||||
std::vector<ov::test::InputShape> expandedShapes(10, inputShape);
|
||||
std::vector<ov::test::InputShape> expandedShapes(9, inputShape);
|
||||
init_input_shapes(expandedShapes);
|
||||
|
||||
auto f = ov::test::snippets::EltwiseMaxNumParamsFunction(inputDynamicShapes);
|
||||
|
|
|
|||
|
|
@ -94,18 +94,18 @@ public:
|
|||
protected:
|
||||
std::shared_ptr<ov::Model> initOriginal() const override;
|
||||
};
|
||||
/// Eltwise graph with 10 inputs and 2 outputs.
|
||||
/// Eltwise graph with 9 inputs and 2 outputs.
|
||||
/// Needed to test for a max number of inputs+outputs allowed.
|
||||
// in1 in2 in3 ... in10
|
||||
// in1 in2 in3 ... in9
|
||||
// ........................
|
||||
// Subtract Power
|
||||
// \ Sinh
|
||||
// Subtract Power
|
||||
// \ Sinh
|
||||
// Result
|
||||
class EltwiseMaxNumParamsFunction : public SnippetsFunctionBase {
|
||||
public:
|
||||
explicit EltwiseMaxNumParamsFunction(const std::vector<PartialShape>& inputShapes) :
|
||||
SnippetsFunctionBase(inputShapes) {
|
||||
OPENVINO_ASSERT(input_shapes.size() == 10, "Got invalid number of input shapes");
|
||||
OPENVINO_ASSERT(input_shapes.size() == 9, "Got invalid number of input shapes");
|
||||
}
|
||||
protected:
|
||||
std::shared_ptr<ov::Model> initOriginal() const override;
|
||||
|
|
|
|||
|
|
@ -91,7 +91,7 @@ std::shared_ptr<ov::Model> EltwiseMaxNumParamsFunction::initOriginal() const {
|
|||
auto param = std::make_shared<op::v0::Parameter>(precision, shape);
|
||||
params.push_back(param);
|
||||
}
|
||||
std::vector<std::shared_ptr<Node>> add; // 5
|
||||
std::vector<std::shared_ptr<Node>> add; // 4
|
||||
for (size_t i = 0; i < input_shapes.size() / 2; i++) {
|
||||
add.push_back(std::make_shared<op::v1::Add>(params[i * 2], params[i * 2 + 1]));
|
||||
}
|
||||
|
|
@ -101,7 +101,7 @@ std::shared_ptr<ov::Model> EltwiseMaxNumParamsFunction::initOriginal() const {
|
|||
mul.push_back(mul_node);
|
||||
}
|
||||
auto sub = std::make_shared<op::v1::Subtract>(mul[0], mul[1]);
|
||||
auto power = std::make_shared<op::v1::Power>(add.back(), sub);
|
||||
auto power = std::make_shared<op::v1::Power>(params.back(), sub);
|
||||
auto exit_sinh = std::make_shared<op::v0::Sinh>(power);
|
||||
return std::make_shared<ov::Model>(NodeVector{sub, exit_sinh}, params);
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue