Cleanup

Author: v-Golubev

src/common/snippets/include/snippets/lowered/pass/brgemm_blocking.hpp

@@ -28,11 +28,10 @@ class BrgemmBlockingBase : public snippets::lowered::pass::RangedPass {
     /**
      * @interface mark_blocking_loops
      * @brief Covers brgemm with blocking loops. Also should calculate optimal blocking parameters inside.
-     * @param linear_ir LIR that contain's brgemm
+     * @param linear_ir LIR that contains brgemm
      * @param brgemm_it iterator on brgemm expression which should be covered with blocking loops
      */
     virtual bool mark_blocking_loops(snippets::lowered::LinearIR& linear_ir, const snippets::lowered::LinearIR::constExprIt& brgemm_it) = 0;
-    // virtual std::tuple<size_t, size_t, size_t> get_blocking_params(const ov::snippets::lowered::ExpressionPtr& brgemm_expr) = 0;
 
     static bool blocking_loop_exists(const snippets::lowered::LoopManagerPtr& loop_manager,
                                      const ov::snippets::lowered::ExpressionPtr& brgemm_expr,

src/common/snippets/include/snippets/pass/explicit_transpose_matmul_inputs.hpp

@@ -30,7 +30,7 @@ class ExplicitTransposeMatMulInputs: public ov::pass::MatcherPass {
     static bool are_weights_scalar(const std::shared_ptr<ov::Node>& node);
 
 private:
-    static bool extract(const ov::Input<ov::Node>& input);
+    static void extract(const ov::Input<ov::Node>& input);
 };
 
 }  // namespace pass

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

@@ -19,7 +19,7 @@ bool ov::snippets::pass::ExplicitTransposeMatMulInputs::are_weights_scalar(const
                        });
 }
 
-bool ov::snippets::pass::ExplicitTransposeMatMulInputs::extract(const ov::Input<ov::Node>& input) {
+void ov::snippets::pass::ExplicitTransposeMatMulInputs::extract(const ov::Input<ov::Node>& input) {
     auto parent = input.get_source_output().get_node_shared_ptr();
     auto transpose = ov::as_type_ptr<ov::op::v1::Transpose>(parent);
     while (!transpose && !ov::is_type<ov::op::v0::Parameter>(parent)) {
@@ -47,7 +47,7 @@ bool ov::snippets::pass::ExplicitTransposeMatMulInputs::extract(const ov::Input<
         new_transpose_order->set_friendly_name(transpose_pattern->get_friendly_name());
         ov::copy_runtime_info(transpose_pattern, new_transpose_order);
         transpose->set_argument(1, new_transpose_order);
-        return true;
+        return;
     }
 
     // Create new Transpose before Parameter
@@ -68,7 +68,6 @@ bool ov::snippets::pass::ExplicitTransposeMatMulInputs::extract(const ov::Input<
     const auto new_transpose = std::make_shared<opset1::Transpose>(parent, constant_order); // parent is Parameter
     const auto consumer_input = *(consumers.begin());
     consumer_input.replace_source_output(new_transpose);
-    return true;
 }
 
 ov::snippets::pass::ExplicitTransposeMatMulInputs::ExplicitTransposeMatMulInputs() {
@@ -86,12 +85,14 @@ ov::snippets::pass::ExplicitTransposeMatMulInputs::ExplicitTransposeMatMulInputs
             if (!matmul)
                 return false;
 
-            if (matmul->get_transpose_a() && extract(matmul->input(0))) {
+            if (matmul->get_transpose_a()) {
+                extract(matmul->input(0));
                 matmul->set_transpose_a(false);
                 rewritten |= true;
             }
 
-            if (matmul->get_transpose_b() && !transformation_callback(matmul) && extract(matmul->input(1))) {
+            if (matmul->get_transpose_b() && !transformation_callback(matmul)) {
+                extract(matmul->input(1));
                 matmul->set_transpose_b(false);
                 rewritten |= true;
             }

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

@@ -73,8 +73,9 @@ size_t jit_brgemm_copy_b_emitter::compute_vnni_factor(const ov::element::Type& p
 }
 
 size_t jit_brgemm_copy_b_emitter::get_elems_in_vec(const ov::element::Type& precision) {
-    OV_CPU_JIT_EMITTER_ASSERT(dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core), "doesn't support non avx512 platforms");
-    const auto vlen = dnnl::impl::cpu::x64::cpu_isa_traits<dnnl::impl::cpu::x64::avx512_core>::vlen;
+    using namespace dnnl::impl::cpu;
+    OV_CPU_JIT_EMITTER_ASSERT(x64::mayiuse(x64::avx2), "doesn't support non avx512 platforms");
+    const auto vlen = x64::mayiuse(avx512_core) ? x64::cpu_isa_traits<x64::avx512_core>::vlen : x64::cpu_isa_traits<x64::avx2>::vlen;
     return vlen / precision.size();
 }
 
@@ -84,6 +85,7 @@ jit_brgemm_copy_b_emitter::jit_brgemm_copy_b_emitter(jit_generator* h, cpu_isa_t
     const auto brgemm_repack = ov::as_type_ptr<ov::intel_cpu::BrgemmCopyB>(expr->get_node());
     if (!brgemm_repack)
         OV_CPU_JIT_EMITTER_THROW("expects BrgemmCopyB node");
+    OV_CPU_JIT_EMITTER_ASSERT(is_superset(host_isa_, cpu::x64::avx2), "host_isa must be at least avx2");
     m_with_comp = with_compensations(brgemm_repack->get_type());
     m_in_offset = brgemm_repack->get_offset_in();
     m_out_offset = brgemm_repack->get_offset_out();

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

@@ -25,12 +25,36 @@ class jit_brgemm_copy_b_emitter : public jit_emitter {
         return {{element::i8}, {element::bf16}, {element::f32}};
     }
 
+    /**
+     * @brief Computes buffer size that OneDNN impl needs for repacked tensor
+     * @param copy_b_expr Repacking expression whose information (tensor precision, layout, subtensors) is used for
+     * buffer size computations
+     */
     static size_t get_repacking_buffer_size(const ov::snippets::lowered::ExpressionPtr& copy_b_expr);
+    /**
+     * @brief Computes buffer size that OneDNN impl needs for compensations
+     * @param copy_b_expr Repacking expression whose information (tensor precision, subtensors) is used for
+     * buffer size computations
+     */
     static size_t get_compensations_buffer_size(const ov::snippets::lowered::ExpressionPtr& copy_b_expr);
 
+    /**
+     * @brief Computes leading dimension (LDB) which must be used in brgemm and brgemm_copy_b emitters
+     * @param n_block N block size shared between BrgemmCPU and BrgemmCopyB node
+     * @param precision tensor precision
+     */
     static size_t compute_repacking_out_leading_dim(const size_t n_block, const ov::element::Type& precision);
+    /**
+     * @brief Computes inner N block size used by OneDNN implementation. Depends on tensor precision
+     */
     static size_t compute_inner_n_block(const ov::element::Type& precision);
+    /**
+     * @brief Computes VNNI factor used by OneDNN implementation. Depends on tensor precision
+     */
     static size_t compute_vnni_factor(const ov::element::Type& precision);
+    /**
+     * @brief Computes number of elems with requested precision that fit in the vector register
+     */
     static size_t get_elems_in_vec(const ov::element::Type& precision);
 
 private:

src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_cpu_blocking.cpp

@@ -13,7 +13,6 @@
 #include "snippets/utils/utils.hpp"
 #include "transformations/snippets/x64/op/brgemm_cpu.hpp"
 #include "transformations/snippets/x64/op/brgemm_utils.hpp"
-#include "transformations/tpp/x64/op/brgemm.hpp"
 
 
 namespace ov {
@@ -23,6 +22,7 @@ using LinearIR = snippets::lowered::LinearIR;
 using LoopPort = snippets::lowered::LoopPort;
 using ExpressionPtr = ov::snippets::lowered::ExpressionPtr;
 using namespace ov::snippets::lowered;
+using namespace ov::intel_cpu::brgemm_utils;
 
 LinearIR::constExprIt BrgemmCPUBlocking::move_new_memory_buffer(LinearIR& linear_ir, const LinearIR::constExprIt& brgemm_it) {
     const auto& brgemm_expr = brgemm_it->get();
@@ -42,7 +42,7 @@ LinearIR::constExprIt BrgemmCPUBlocking::get_loop_begin_pos(LinearIR& linear_ir,
     const auto& brgemm_expr = *brgemm_it;
     const auto brgemm = ov::as_type_ptr<intel_cpu::BrgemmCPU>(brgemm_expr->get_node());
     OPENVINO_ASSERT(brgemm, "get_loop_begin_pos must be called only for BrgemmCPU expression");
-    if (ov::intel_cpu::brgemm_utils::with_amx(brgemm->get_type()))
+    if (with_amx(brgemm->get_type()))
         loop_begin_it = move_new_memory_buffer(linear_ir, brgemm_it);
     if (copy_b_expr)
         loop_begin_it = linear_ir.find(copy_b_expr);
@@ -72,7 +72,6 @@ bool BrgemmCPUBlocking::mark_blocking_loops(LinearIR& linear_ir, const LinearIR:
     const auto& k = *in_0_planar_dims.rbegin();
     OPENVINO_ASSERT(k == *++in_1_planar_dims.rbegin(), "Brgemm input descriptors have different K dimension value.");
     const auto type = brgemm->get_type();
-    const bool with_repacking = ov::intel_cpu::brgemm_utils::with_repacking(type);
 
     // Ticket: 113745
     // TODO: extend block size selection heuristics
@@ -83,15 +82,15 @@ bool BrgemmCPUBlocking::mark_blocking_loops(LinearIR& linear_ir, const LinearIR:
         // K blocking is disabled in dynamism by default
         if (ov::snippets::utils::is_dynamic_value(K))
             return snippets::utils::get_dynamic_value<size_t>();
-        if (with_repacking)
+        if (with_repacking(type))
             return K;
         return K > 1024 ? 1024 : K > 512 ? 512 : K;
     };
     auto get_block_size_n = [&](const size_t N) -> size_t {
         // N blocking is disabled in dynamism by default
         if (ov::snippets::utils::is_dynamic_value(N))
             return snippets::utils::get_dynamic_value<size_t>();
-        if (with_repacking)
+        if (with_repacking(type))
             return N;
         return std::min<size_t>(64, N);
     };
@@ -105,13 +104,13 @@ bool BrgemmCPUBlocking::mark_blocking_loops(LinearIR& linear_ir, const LinearIR:
     brgemm_expr->get_output_port_descriptor(0)->set_subtensor(ov::snippets::VectorDims{block_size_m, block_size_n});
 
     ov::snippets::lowered::ExpressionPtr copy_b_expr = nullptr;
-    if (with_repacking) {
+    if (with_repacking(type)) {
         const auto copy_b = brgemm->get_brgemm_copy();
         copy_b_expr = linear_ir.get_expr_by_node(copy_b);
         const ov::snippets::VectorDims repacking_subtensor{block_size_k, block_size_n};
         copy_b_expr->get_input_port_descriptor(0)->set_subtensor(repacking_subtensor);
         copy_b_expr->get_output_port_descriptor(0)->set_subtensor(repacking_subtensor);
-        if (ov::intel_cpu::brgemm_utils::with_compensations(type)) {
+        if (with_compensations(type)) {
             const ov::snippets::VectorDims compensations_subtensor{1, block_size_n};
             OPENVINO_ASSERT(brgemm_expr->get_input_count() == 3, "Brgemm must have 3 inputs in case of compensations.");
             brgemm_expr->get_input_port_descriptor(2)->set_subtensor(compensations_subtensor);
@@ -126,7 +125,7 @@ bool BrgemmCPUBlocking::mark_blocking_loops(LinearIR& linear_ir, const LinearIR:
 
         const auto b_input_port = include_repacking && copy_b_expr ? copy_b_expr->get_input_port(0) : brgemm_expr->get_input_port(1);
         std::vector<LoopPort> entries{LoopPort(brgemm_expr->get_input_port(0), true), LoopPort(b_input_port, false)};
-        if (!include_repacking && ov::intel_cpu::brgemm_utils::with_compensations(type))
+        if (!include_repacking && with_compensations(type))
             entries.emplace_back(brgemm_expr->get_input_port(2), 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);

src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp

@@ -1105,8 +1105,7 @@ void Transformations::MainSnippets(void) {
     },
     snippets::pass::TokenizeSnippets);
 
-    CPU_SET_CALLBACK_COMMON(snippetsManager,
-    [this](const std::shared_ptr<const ov::Node>& n) -> bool {
+    auto mm_supports_transpose_b = [this](const std::shared_ptr<const ov::Node>& n) {
         MAYBE_UNUSED(inferencePrecision);
         const auto& b_shape = n->get_input_partial_shape(1);
         if (b_shape.is_dynamic())
@@ -1137,6 +1136,11 @@ void Transformations::MainSnippets(void) {
             return false;
 #endif
         return true;
+    };
+
+    CPU_SET_CALLBACK_COMMON(snippetsManager,
+    [&mm_supports_transpose_b](const std::shared_ptr<const ov::Node>& n) {
+        return mm_supports_transpose_b(n);
     },
     snippets::pass::ExplicitTransposeMatMulInputs);