[CPU][ARM]Snippets MatMul via brgemm emitter and executor #28304
+689
−83
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4,7 +4,7 @@ | |
|
|
||
| #pragma once | ||
|
|
||
| #include "emitters/snippets/brgemm_base.hpp" | ||
|
|
||
| namespace ov { | ||
| namespace intel_cpu { | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,85 @@ | ||
| // Copyright (C) 2025 Intel Corporation | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
| // | ||
|
|
||
| #include "jit_brgemm_emitter.hpp" | ||
|
|
||
| #include "snippets/utils/utils.hpp" | ||
| #include "transformations/tpp/common/op/brgemm.hpp" | ||
|
|
||
| using namespace Xbyak_aarch64; | ||
|
|
||
| namespace ov { | ||
| namespace intel_cpu { | ||
| namespace aarch64 { | ||
|
|
||
| using jit_generator = dnnl::impl::cpu::aarch64::jit_generator; | ||
| using cpu_isa_t = dnnl::impl::cpu::aarch64::cpu_isa_t; | ||
| using ExpressionPtr = ov::snippets::lowered::ExpressionPtr; | ||
|
|
||
| jit_brgemm_emitter::jit_brgemm_emitter(jit_generator* h, | ||
| cpu_isa_t isa, | ||
| const ExpressionPtr& expr, | ||
| const snippets::KernelExecutorTablePtr& kernel_table, | ||
| const ov::intel_cpu::MultiCacheWeakPtr& compiled_kernel_cache) | ||
| : jit_emitter(h, isa) { | ||
| in_out_type_ = emitter_in_out_map::gpr_to_gpr; | ||
| const auto& brgemm_node = as_type_ptr<intel_cpu::tpp::op::BrgemmTPP>(expr->get_node()); | ||
| const auto& brg0Prc = brgemm_node->get_input_element_type(0); | ||
| const auto& brg1Prc = brgemm_node->get_input_element_type(1); | ||
| BrgemmKernelConfig kernel_config(brg0Prc, brg1Prc, isa); | ||
| m_kernel_executor = kernel_table->register_kernel<BrgemmKernelExecutor>(expr, compiled_kernel_cache, kernel_config); | ||
| } | ||
|
|
||
| std::set<std::vector<element::Type>> jit_brgemm_emitter::get_supported_precisions( | ||
| const std::shared_ptr<ov::Node>& node) { | ||
| // Note: Brgemm currently supports only fp32 on arm | ||
| return {{element::f32, element::f32}}; | ||
| } | ||
|
|
||
| void jit_brgemm_emitter::validate_arguments(const std::vector<size_t>& in, const std::vector<size_t>& out) const { | ||
| OV_CPU_JIT_EMITTER_ASSERT(in.size() == 2, "Expects 2 input regs, got" + std::to_string(in.size())); | ||
| OV_CPU_JIT_EMITTER_ASSERT(out.size() == 1, "Expects 1 output reg, got" + std::to_string(out.size())); | ||
| } | ||
|
|
||
| void jit_brgemm_emitter::emit_code(const std::vector<size_t>& in, | ||
| const std::vector<size_t>& out, | ||
| const std::vector<size_t>& pool_vec_idxs, | ||
| const std::vector<size_t>& pool_gpr_idxs) const { | ||
| validate_arguments(in, out); | ||
| emit_impl(in, out); | ||
| } | ||
|
|
||
| void jit_brgemm_emitter::emit_impl(const std::vector<size_t>& in, const std::vector<size_t>& out) const { | ||
| validate_arguments(in, out); | ||
| std::unordered_set<size_t> exclude = {}; | ||
| store_context(exclude); | ||
|
There was a problem hiding this comment. Please note that we will merge #27391 soon. This PR efficently provides efficient work with reg spills - we will able to spill only needed (live) registers Just for information and to align with other our activities 😊 |
||
|
|
||
| Xbyak_aarch64::XReg func_reg(9); | ||
| h->mov(func_reg, get_execute_function_ptr()); | ||
| Xbyak_aarch64::XReg x0(0); | ||
| Xbyak_aarch64::XReg x1(1); | ||
| Xbyak_aarch64::XReg x2(2); | ||
| Xbyak_aarch64::XReg x3(3); | ||
|
|
||
| const auto& compiled_kernel = get_compiled_kernel_ptr(); | ||
| h->mov(x0, compiled_kernel); | ||
| h->mov(x1, Xbyak_aarch64::XReg(in[0])); | ||
| h->mov(x2, Xbyak_aarch64::XReg(in[1])); | ||
| h->mov(x3, Xbyak_aarch64::XReg(out[0])); | ||
| h->blr(func_reg); | ||
|
|
||
| restore_context(exclude); | ||
| } | ||
|
|
||
| const uintptr_t jit_brgemm_emitter::get_compiled_kernel_ptr() const { | ||
| return reinterpret_cast<const uintptr_t>(m_kernel_executor.get()); | ||
| } | ||
|
|
||
| const uintptr_t jit_brgemm_emitter::get_execute_function_ptr() const { | ||
| return reinterpret_cast<const uintptr_t>(BrgemmKernelExecutor::execute); | ||
| } | ||
|
|
||
| } // namespace aarch64 | ||
| } // namespace intel_cpu | ||
| } // namespace ov | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,46 @@ | ||
| // Copyright (C) 2025 Intel Corporation | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
| // | ||
|
|
||
| #pragma once | ||
|
|
||
| #include "emitters/plugin/aarch64/jit_emitter.hpp" | ||
| #include "emitters/tpp/aarch64/kernel_executors/brgemm.hpp" | ||
|
|
||
| namespace ov { | ||
| namespace intel_cpu { | ||
| namespace aarch64 { | ||
|
|
||
| class jit_brgemm_emitter : public jit_emitter { | ||
| public: | ||
| jit_brgemm_emitter(dnnl::impl::cpu::aarch64::jit_generator* h, | ||
| dnnl::impl::cpu::aarch64::cpu_isa_t isa, | ||
| const ov::snippets::lowered::ExpressionPtr& expr, | ||
| const snippets::KernelExecutorTablePtr& kernel_table, | ||
| const ov::intel_cpu::MultiCacheWeakPtr& compiled_kernel_cache); | ||
|
|
||
| size_t get_inputs_count() const override { | ||
| return 2; | ||
| } | ||
|
|
||
| static std::set<std::vector<element::Type>> get_supported_precisions( | ||
| const std::shared_ptr<ov::Node>& node = nullptr); | ||
|
|
||
| void emit_code(const std::vector<size_t>& in_idxs, | ||
| const std::vector<size_t>& out_idxs, | ||
| const std::vector<size_t>& pool_vec_idxs = {}, | ||
| const std::vector<size_t>& pool_gpr_idxs = {}) const override; | ||
|
|
||
| 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; | ||
|
|
||
| const uintptr_t get_execute_function_ptr() const; | ||
| const uintptr_t get_compiled_kernel_ptr() const; | ||
|
|
||
| std::shared_ptr<BrgemmKernelExecutor> m_kernel_executor = nullptr; | ||
| }; | ||
|
|
||
| } // namespace aarch64 | ||
| } // namespace intel_cpu | ||
| } // namespace ov |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There is also RISCV64, AArch32 etc. Can we add only supported archs to condition?