[Snippets] Support FP32/BF16/I8 matmuls with transpose_b=true via BrgemmCopyB

Author: v-Golubev

src/common/snippets/docs/mha_optimization_guide.md

@@ -123,7 +123,7 @@ For enhancing the execution efficiency, blocking across the M, K, and N matmul d
 
 ### Blocking Parameters
 
-The heuristics for determining the optimal block sizes can be found in [SetBrgemmCPUBlockingParams](../../../plugins/intel_cpu/src/transformations/snippets/x64/pass/set_brgemm_cpu_blocking_params.cpp).
+The heuristics for determining the optimal block sizes can be found in [BrgemmCPUBlocking](../../../plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_cpu_blocking.cpp).
 
 **Please note: Blocking by M dimension is shared between both Brgemms. Please see [SplitLoops](../include/snippets/lowered/pass/split_loops.hpp) lowered pass for the details.**
 
@@ -141,7 +141,7 @@ Based on previously discussed information, we provide the following recommendati
 In local experiments, some transformations might be worth to change:
     - Disable [ExtractUnsupportedTransposes](#extractunsupportedtransposes) transformation in order to benchmark Snippets Transpose implementation.
     - Adjust [SplitDimensionM](#splitdimensionm) heuristics in order to benchmark another splitting, or disable the pass at all.
-3. [Blocking parameters](#blocking-parameters): adjust blocking heuristics in `SetBrgemmCPUBlockingParams`.
+3. [Blocking parameters](#blocking-parameters): adjust blocking heuristics in `BrgemmCPUBlocking`.
     - Please note that there are 2 Matmul nodes inside a single MHA, and each Matmul can have his own optimal K, N blocking params.
     M block is better to keep the same since the corresponding blocking loop is shared between both Matmuls.
     - For the BF16/INT8 blocking loops, 2 options are possible: blocking can be done only for Brgemm node, or for BrgemmCopyB repacking too.

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

@@ -0,0 +1,47 @@
+// Copyright (C) 2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "snippets/lowered/pass/pass.hpp"
+#include "snippets/lowered/specific_loop_iter_handlers.hpp"
+#include "snippets/op/brgemm.hpp"
+
+namespace ov {
+namespace snippets {
+namespace lowered {
+namespace pass {
+
+/**
+ * @interface BrgemmBlockingBase
+ * @brief Base class for Brgemm blocking loops markup
+ * @ingroup snippets
+ */
+class BrgemmBlockingBase : public snippets::lowered::pass::RangedPass {
+public:
+    OPENVINO_RTTI("BrgemmBlockingBase", "RangedPass")
+    bool run(snippets::lowered::LinearIR& linear_ir,
+             snippets::lowered::LinearIR::constExprIt begin,
+             snippets::lowered::LinearIR::constExprIt end) override;
+
+    static snippets::lowered::SpecificIterationHandlers get_default_blocking_loop_handlers(size_t work_amount, size_t block_size);
+
+protected:
+    /**
+     * @interface mark_blocking_loops
+     * @brief Covers brgemm with blocking loops. Also should calculate optimal blocking parameters inside.
+     * @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;
+
+    static bool blocking_loop_exists(const snippets::lowered::LoopManagerPtr& loop_manager,
+                                     const ov::snippets::lowered::ExpressionPtr& brgemm_expr,
+                                     const std::shared_ptr<ov::snippets::op::Brgemm>& brgemm);
+};
+
+} // namespace pass
+} // namespace lowered
+} // namespace snippets
+} // namespace ov

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

@@ -64,6 +64,21 @@ class TransformInnerSplitLoop : public pass::RangedPass {
     size_t m_tail_size;
 };
 
+/**
+ * @interface SetEvaluateOnce
+ * @brief The pass set `evaluate once = true` only to ExpandedLoopInfo which is mapped on LoopEnd in the passed iterator `end`.
+ *        The pointer arithmetic should be updated in the separate optimization `OptimizeLoopSingleEvaluation`
+ * @ingroup snippets
+ */
+class SetEvaluateOnce : public snippets::lowered::pass::RangedPass {
+public:
+    SetEvaluateOnce() = default;
+    OPENVINO_RTTI("SetEvaluateOnce", "RangedPass")
+    bool run(snippets::lowered::LinearIR& linear_ir,
+             snippets::lowered::LinearIR::constExprIt begin,
+             snippets::lowered::LinearIR::constExprIt end) override;
+    std::shared_ptr<snippets::lowered::pass::PassBase> merge(const std::shared_ptr<snippets::lowered::pass::PassBase>& other) override;
+};
 } // namespace pass
 } // namespace lowered
 } // namespace snippets

src/common/snippets/include/snippets/op/brgemm.hpp

@@ -22,26 +22,17 @@ class Brgemm : virtual public modifier::MemoryAccess, public ov::op::Op {
     OPENVINO_OP("Brgemm", "SnippetsOpset");
     Brgemm(const Output<Node>& A, const Output<Node>& B,
            const size_t offset_a = 0lu, const size_t offset_b = 0lu, const size_t offset_c = 0lu,
-           std::vector<size_t> layout_a = {}, std::vector<size_t> layout_b = {}, std::vector<size_t> layout_c = {},
-           size_t blk_size_m = 0, size_t blk_size_k = 0, size_t blk_size_n = 0);
+           std::vector<size_t> layout_a = {}, std::vector<size_t> layout_b = {}, std::vector<size_t> layout_c = {});
     Brgemm(const Output<Node>& A, const Output<Node>& B,
            const PortDescriptor& desc_a, const PortDescriptor& desc_b, const PortDescriptor& desc_c,
-           std::vector<size_t> layout_a = {}, std::vector<size_t> layout_b = {}, std::vector<size_t> layout_c = {},
-           size_t blk_size_m = 0, size_t blk_size_k = 0, size_t blk_size_n = 0);
+           std::vector<size_t> layout_a = {}, std::vector<size_t> layout_b = {}, std::vector<size_t> layout_c = {});
     Brgemm() = default;
 
     size_t get_offset_a() const { return get_input_offset(0); }
     size_t get_offset_b() const { return get_input_offset(1); }
     size_t get_offset_c() const { return get_output_offset(0); }
 
-    size_t get_m_block_size() const { return m_M_blk; }
-    size_t get_k_block_size() const { return m_K_blk; }
-    size_t get_n_block_size() const { return m_N_blk; }
     float get_beta() const { return m_beta; }
-
-    void set_m_block_size(size_t block_size) { m_M_blk = block_size; }
-    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; }
     void set_beta(float beta) { m_beta = beta; }
 
     static ov::element::Type get_output_type(const ov::element::Type& in_type0, const ov::element::Type& in_type1);
@@ -57,10 +48,6 @@ class Brgemm : virtual public modifier::MemoryAccess, public ov::op::Op {
     std::vector<ov::PartialShape> get_planar_input_shapes(const std::vector<ov::Input<ov::Node>>& inputs) const;
     ov::PartialShape infer_output_partial_shape(const std::vector<ov::PartialShape>& input_shapes) const;
     ov::PartialShape get_planar_output_shape(const ov::PartialShape& output_shape) const;
-    void set_block_size_values(size_t blk_size_m, size_t blk_size_k, size_t blk_size_n);
-    size_t m_M_blk = 0;
-    size_t m_K_blk = 0;
-    size_t m_N_blk = 0;
     float m_beta = 0.f;
 
 private:

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

@@ -22,9 +22,6 @@ class MatMulToBrgemm: public ov::pass::MatcherPass {
 public:
     OPENVINO_RTTI("MatMulToBrgemm", "0");
     MatMulToBrgemm();
-
-private:
-    void init_ports(const std::shared_ptr<op::Brgemm>& brgemm) const;
 };
 
 

src/common/snippets/include/snippets/utils/utils.hpp

@@ -127,6 +127,10 @@ bool broadcast_merge_dim(size_t& dst, const size_t& d1, const size_t& d2);
 VectorDims pshape_to_vdims(const PartialShape&);
 ov::PartialShape vdims_to_pshape(const VectorDims&);
 
+inline size_t dimension_to_size_t(const ov::Dimension& dim) {
+    return dim.is_dynamic() ? snippets::utils::get_dynamic_value<VectorDims::value_type>() : static_cast<size_t>(dim.get_length());
+}
+
 // dim_idx starts from the layout end: dim_idx = 0 -> last element in layout (layout.back())
 inline size_t get_input_dim_idx(const std::vector<size_t>& layout, size_t dim_idx) {
     OPENVINO_ASSERT(dim_idx < layout.size(), "Incorrect dim_idx");

src/common/snippets/src/lowered/pass/brgemm_blocking.cpp

@@ -0,0 +1,67 @@
+// Copyright (C) 2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "snippets/lowered/pass/brgemm_blocking.hpp"
+
+#include "snippets/itt.hpp"
+#include "snippets/lowered/linear_ir.hpp"
+#include "snippets/lowered/loop_manager.hpp"
+#include "snippets/lowered/pass/pass.hpp"
+#include "snippets/lowered/pass/propagate_subtensors.hpp"
+#include "snippets/lowered/pass/iter_handler.hpp"
+#include "snippets/snippets_isa.hpp"
+#include "snippets/utils/utils.hpp"
+
+namespace ov {
+namespace snippets {
+namespace lowered {
+namespace pass {
+
+snippets::lowered::SpecificIterationHandlers BrgemmBlockingBase::get_default_blocking_loop_handlers(size_t work_amount, size_t block_size) {
+    SpecificIterationHandlers handlers;
+    const auto tail_size = snippets::utils::is_dynamic_value(work_amount) ? snippets::utils::get_dynamic_value<size_t>() : work_amount % block_size;
+    if (tail_size != 0)
+        handlers.register_pass<snippets::lowered::SpecificLoopIterType::LAST_ITER, snippets::lowered::pass::UpdateSubtensors>(tail_size);
+    handlers.register_pass<snippets::lowered::SpecificLoopIterType::LAST_ITER, SetEvaluateOnce>();
+    return handlers;
+}
+
+bool BrgemmBlockingBase::blocking_loop_exists(const snippets::lowered::LoopManagerPtr& loop_manager,
+                                              const ExpressionPtr& brgemm_expr,
+                                              const std::shared_ptr<snippets::op::Brgemm>& brgemm) {
+    auto check_port = [&](const LoopPort& p) {
+        return p.expr_port->get_expr() == brgemm_expr && ov::snippets::utils::one_of(p.dim_idx, 0ul, 1ul);
+    };
+
+    const auto& loop_ids = brgemm_expr->get_loop_ids();
+    for (const auto& id : loop_ids) {
+        const auto loop = loop_manager->get_loop_info(id);
+        if (std::any_of(loop->get_input_ports().begin(), loop->get_input_ports().end(), check_port) ||
+            std::any_of(loop->get_output_ports().begin(), loop->get_output_ports().end(), check_port)) {
+            return true;
+        }
+    }
+    return false;
+}
+
+bool BrgemmBlockingBase::run(LinearIR& linear_ir, LinearIR::constExprIt begin, LinearIR::constExprIt end) {
+    OV_ITT_SCOPED_TASK(ov::pass::itt::domains::SnippetsTransform, "Snippets::BrgemmCPUBlocking")
+    const auto& loop_manager = linear_ir.get_loop_manager();
+    bool modified = false;
+    for (auto expr_it = begin; expr_it != end; expr_it++) {
+        const auto& brgemm_expr = *expr_it;
+        const auto& node = brgemm_expr->get_node();
+        const auto brgemm = ov::as_type_ptr<ov::snippets::op::Brgemm>(node);
+        if (!brgemm || blocking_loop_exists(loop_manager, brgemm_expr, brgemm))
+            continue;
+        modified = mark_blocking_loops(linear_ir, expr_it);
+    }
+
+    return modified;
+}
+
+} // namespace pass
+} // namespace lowered
+} // namespace snippets
+} // namespace ov

src/common/snippets/src/lowered/pass/iter_handler.cpp

@@ -142,6 +142,18 @@ std::shared_ptr<pass::PassBase> TransformInnerSplitLoop::merge(const std::shared
     return merged_pass;
 }
 
+bool SetEvaluateOnce::run(LinearIR& linear_ir, LinearIR::constExprIt begin, LinearIR::constExprIt end) {
+    const auto& loop_end = ov::as_type_ptr<snippets::op::LoopEnd>(end->get()->get_node());
+    OPENVINO_ASSERT(loop_end, "SetEvaluateOnce expected LoopEnd node in iterator `end`.");
+    const auto& loop_info = linear_ir.get_loop_manager()->get_loop_info<ov::snippets::lowered::ExpandedLoopInfo>(loop_end->get_id());
+    loop_info->set_evaluate_once(true);
+    return true;
+}
+
+std::shared_ptr<snippets::lowered::pass::PassBase> SetEvaluateOnce::merge(const std::shared_ptr<snippets::lowered::pass::PassBase>& other) {
+    return !other || ov::is_type<SetEvaluateOnce>(other) ? std::make_shared<SetEvaluateOnce>() : nullptr;
+}
+
 } // namespace pass
 } // namespace lowered
 } // namespace snippets

src/common/snippets/src/op/brgemm.cpp

@@ -32,33 +32,23 @@ std::vector<size_t> get_output_layout(const std::shared_ptr<const ov::Node>& n)
 
 Brgemm::Brgemm(const Output<Node>& A, const Output<Node>& B,
                const size_t offset_a, const size_t offset_b, const size_t offset_c,
-               std::vector<size_t> layout_a, std::vector<size_t> layout_b, std::vector<size_t> layout_c,
-               const size_t blk_size_m, const size_t blk_size_k, const size_t blk_size_n)
+               std::vector<size_t> layout_a, std::vector<size_t> layout_b, std::vector<size_t> layout_c)
     : MemoryAccess(std::set<size_t>{0, 1}, std::set<size_t>{0}), Op({A, B}) {
     set_output_size(1);
     set_input_offset(offset_a, 0);
     set_input_offset(offset_b, 1);
     set_output_offset(offset_c, 0);
-    set_block_size_values(blk_size_m, blk_size_k, blk_size_n);
     custom_constructor_validate_and_infer_types(std::move(layout_a), std::move(layout_b), std::move(layout_c));
 }
 
 Brgemm::Brgemm(const Output<Node>& A, const Output<Node>& B,
                const PortDescriptor& desc_a, const PortDescriptor& desc_b, const PortDescriptor& desc_c,
-               std::vector<size_t> layout_a, std::vector<size_t> layout_b, std::vector<size_t> layout_c,
-               const size_t blk_size_m, const size_t blk_size_k, const size_t blk_size_n)
+               std::vector<size_t> layout_a, std::vector<size_t> layout_b, std::vector<size_t> layout_c)
     : MemoryAccess(PortMap{{0, desc_a}, {1, desc_b}}, PortMap{{0, desc_c}}), Op({A, B}) {
     set_output_size(1);
-    set_block_size_values(blk_size_m, blk_size_k, blk_size_n);
     custom_constructor_validate_and_infer_types(std::move(layout_a), std::move(layout_b), std::move(layout_c));
 }
 
-void Brgemm::set_block_size_values(const size_t blk_size_m, const size_t blk_size_k, const size_t blk_size_n) {
-    m_M_blk = blk_size_m;
-    m_K_blk = blk_size_k;
-    m_N_blk = blk_size_n;
-}
-
 void Brgemm::custom_constructor_validate_and_infer_types(std::vector<size_t> layout_a, std::vector<size_t> layout_b, std::vector<size_t> layout_c) {
     INTERNAL_OP_SCOPE(BrgemmCPU_constructor_validate_and_infer_types);
 
@@ -90,9 +80,6 @@ std::shared_ptr<Node> Brgemm::clone_with_new_inputs(const OutputVector& new_args
 }
 
 bool Brgemm::visit_attributes(AttributeVisitor& visitor) {
-    visitor.on_attribute("blk_M", m_M_blk);
-    visitor.on_attribute("blk_K", m_K_blk);
-    visitor.on_attribute("blk_N", m_N_blk);
     visitor.on_attribute("beta", m_beta);
     return MemoryAccess::visit_attributes(visitor);
 }

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

@@ -4,12 +4,12 @@
 
 #include "snippets/pass/explicit_transpose_matmul_inputs.hpp"
 
-#include "snippets/op/subgraph.hpp"
-#include "snippets/itt.hpp"
-
+#include "openvino/core/rt_info.hpp"
 #include "openvino/pass/pattern/matcher.hpp"
 #include "openvino/pass/pattern/op/wrap_type.hpp"
-#include "openvino/core/rt_info.hpp"
+#include "snippets/itt.hpp"
+#include "snippets/op/subgraph.hpp"
+#include "snippets/pass/mha_tokenization.hpp"
 
 bool ov::snippets::pass::ExplicitTransposeMatMulInputs::are_weights_scalar(const std::shared_ptr<ov::Node>& node) {
     const auto inputs = node->inputs();
@@ -58,6 +58,7 @@ void ov::snippets::pass::ExplicitTransposeMatMulInputs::extract(const ov::Input<
                     "ExplicitTransposeMatMulInputs expects Parameter with one consumer in cases when there isn't existing Transpose on input");
     // Extract Transpose from MatMul
     OPENVINO_ASSERT(input.get_partial_shape().rank().is_static(), "ExplicitTransposeMatMulInputs supports only static ranks of shapes");
+
     const auto rank = input.get_partial_shape().size();
     std::vector<size_t> transpose_order(rank, 0);
     std::iota(transpose_order.begin(), transpose_order.end(), 0);
@@ -75,7 +76,7 @@ ov::snippets::pass::ExplicitTransposeMatMulInputs::ExplicitTransposeMatMulInputs
     auto m_matmul0 = std::make_shared<ov::op::v0::MatMul>(ov::pass::pattern::any_input(), ov::pass::pattern::any_input());
 
     register_matcher(std::make_shared<ov::pass::pattern::Matcher>(m_matmul0, matcher_name),
-        [=](ov::pass::pattern::Matcher &m) {
+        [OV_CAPTURE_CPY_AND_THIS](ov::pass::pattern::Matcher &m) {
             OV_ITT_SCOPED_TASK(ov::pass::itt::domains::SnippetsTransform, "Snippets::op::ExplicitTransposeMatMulInputs")
             auto root = m.get_match_root();
             bool rewritten = false;
@@ -89,7 +90,8 @@ ov::snippets::pass::ExplicitTransposeMatMulInputs::ExplicitTransposeMatMulInputs
                 matmul->set_transpose_a(false);
                 rewritten |= true;
             }
-            if (matmul->get_transpose_b()) {
+
+            if (matmul->get_transpose_b() && !transformation_callback(matmul)) {
                 extract(matmul->input(1));
                 matmul->set_transpose_b(false);
                 rewritten |= true;