[Snippets] Support Brgemm with transposed_b via BrgemmCopyB (#24932)

### Details:
 - *Support FP32/BF16/I8 matmuls with transpose_b=true via BrgemmCopyB*
- *BrgemmCopyB emitter: handle tail iteration by N before the main body*
- *Remove workaround on LDB and N dim rounding in brgemm emitters and
related buffers*

### Tickets:
 - *CVS-114487*

## TODO:
- [ ] BufferAllocation test for FP32 brgemm with repacking
- [ ] SetBrgemmCopyBBuffersShape tests
- [ ] MHA with transpose B for low precisions (FP32 already exists)
- [ ] FuseTransposeBrgemm tests

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,107 @@
+// Copyright (C) 2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "snippets/itt.hpp"
+#include "snippets/lowered/loop_manager.hpp"
+#include "snippets/lowered/specific_loop_iter_handlers.hpp"
+#include "snippets/lowered/pass/iter_handler.hpp"
+#include "snippets/op/brgemm.hpp"
+
+namespace ov {
+namespace snippets {
+namespace lowered {
+namespace pass {
+
+/**
+ * @interface BrgemmBlockingBase
+ * @brief Base class for Brgemm blocking, which defines interface for blocking markup,
+ *        and contains default implementation
+ * @ingroup snippets
+ */
+class BrgemmBlockingBase {
+public:
+    static snippets::lowered::SpecificIterationHandlers get_default_blocking_loop_handlers(size_t work_amount, size_t block_size);
+
+protected:
+    /**
+     * @interface get_blocking_params
+     * @brief Computes optimal blocking params for current brgemm expression
+     * @param brgemm_expr Brgemm expression
+     * @return tuple in format (m_block, n_block, k_block)
+     */
+    virtual std::tuple<size_t, size_t, size_t> get_blocking_params(const ov::snippets::lowered::ExpressionPtr& brgemm_expr);
+    /**
+     * @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,
+                                     size_t m_block,
+                                     size_t n_block,
+                                     size_t k_block);
+
+    static bool blocking_loop_exists(const snippets::lowered::LoopManagerPtr& loop_manager,
+                                     const ov::snippets::lowered::ExpressionPtr& brgemm_expr);
+
+    static void mark_m_blocking(const snippets::lowered::LoopManagerPtr& loop_manager,
+                                snippets::lowered::LinearIR::constExprIt loop_begin,
+                                snippets::lowered::LinearIR::constExprIt loop_end,
+                                const std::vector<snippets::lowered::LoopPort>& entries,
+                                const std::vector<snippets::lowered::LoopPort>& exits,
+                                size_t block_size_m);
+
+    static void mark_n_blocking(const snippets::lowered::LoopManagerPtr& loop_manager,
+                                snippets::lowered::LinearIR::constExprIt loop_begin,
+                                snippets::lowered::LinearIR::constExprIt loop_end,
+                                const std::vector<snippets::lowered::LoopPort>& entries,
+                                const std::vector<snippets::lowered::LoopPort>& exits,
+                                size_t block_size_n);
+
+    static void mark_k_blocking(const snippets::lowered::LoopManagerPtr& loop_manager,
+                                snippets::lowered::LinearIR::constExprIt loop_begin,
+                                snippets::lowered::LinearIR::constExprIt loop_end,
+                                const std::vector<snippets::lowered::LoopPort>& entries,
+                                const std::vector<snippets::lowered::LoopPort>& exits,
+                                size_t block_size_k);
+};
+
+/**
+ * @interface BrgemmBlocking
+ * @brief Base class for brgemm blocking passes
+ * @ingroup snippets
+ */
+template <typename BRGEMM_TYPE,
+          typename std::enable_if<std::is_base_of<ov::snippets::op::Brgemm, BRGEMM_TYPE>::value, bool>::type = true>
+class BrgemmBlocking : public snippets::lowered::pass::RangedPass, public BrgemmBlockingBase {
+public:
+    OPENVINO_RTTI("BrgemmBlocking", "RangedPass")
+
+    bool run(snippets::lowered::LinearIR& linear_ir,
+             snippets::lowered::LinearIR::constExprIt begin,
+             snippets::lowered::LinearIR::constExprIt end) override final {  // NOLINT
+        OV_ITT_SCOPED_TASK(ov::pass::itt::domains::SnippetsTransform, "Snippets::BrgemmBlocking")
+        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 brgemm = ov::as_type_ptr<BRGEMM_TYPE>(brgemm_expr->get_node());
+            if (!brgemm)
+                continue;
+            OPENVINO_ASSERT(!blocking_loop_exists(loop_manager, brgemm_expr),
+                            "Brgemm mustn't be covered in loops before blocking pass");
+            size_t m_block, n_block, k_block;
+            std::tie(m_block, n_block, k_block) = get_blocking_params(brgemm_expr);
+            modified = mark_blocking_loops(linear_ir, expr_it, m_block, n_block, k_block);
+        }
+        return modified;
+    }
+};
+} // namespace pass
+} // namespace lowered
+} // namespace snippets
+} // namespace ov

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

@@ -64,6 +64,40 @@ 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;
+};
+
+/**
+ * @interface SetBrgemmBeta
+ * @brief The pass updates all CPUBrgemm nodes with a new beta value
+ * @param m_beta - beta which must be set
+ * @ingroup snippets
+ */
+class SetBrgemmBeta : public snippets::lowered::pass::RangedPass {
+public:
+    SetBrgemmBeta(float beta);
+    OPENVINO_RTTI("SetBrgemmBeta", "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;
+
+private:
+    float m_beta = 0;
+};
 } // namespace pass
 } // namespace lowered
 } // namespace snippets

src/common/snippets/include/snippets/lowered/port_descriptor.hpp

@@ -82,14 +82,14 @@ class PortDescriptor {
 
 class PortDescriptorUtils {
 public:
-    static void set_port_descriptor_ptr(const ov::Input<ov::Node>& n, const PortDescriptorPtr& desc);
-    static void set_port_descriptor_ptr(const ov::Output<ov::Node>& n, const PortDescriptorPtr& desc);
-    static void set_port_descriptor(const ov::Input<ov::Node>& n, std::vector<size_t> subtensor, std::vector<size_t> layout = {});
-    static void set_port_descriptor(const ov::Output<ov::Node>& n, std::vector<size_t> subtensor, std::vector<size_t> layout = {});
+    static void set_port_descriptor_ptr(const ov::Input<ov::Node>& in, const PortDescriptorPtr& desc);
+    static void set_port_descriptor_ptr(const ov::Output<ov::Node>& out, const PortDescriptorPtr& desc);
+    static void set_port_descriptor(const ov::Input<ov::Node>& in, std::vector<size_t> subtensor, std::vector<size_t> layout = {});
+    static void set_port_descriptor(const ov::Output<ov::Node>& out, std::vector<size_t> subtensor, std::vector<size_t> layout = {});
 
     static PortDescriptorPtr get_port_descriptor_ptr(const ov::Input<ov::Node>& in);
-    static PortDescriptorPtr get_port_descriptor_ptr(const ov::Input<const ov::Node>& out);
-    static PortDescriptorPtr get_port_descriptor_ptr(const ov::Output<ov::Node>& in);
+    static PortDescriptorPtr get_port_descriptor_ptr(const ov::Input<const ov::Node>& in);
+    static PortDescriptorPtr get_port_descriptor_ptr(const ov::Output<ov::Node>& out);
     static PortDescriptorPtr get_port_descriptor_ptr(const ov::Output<const ov::Node>& out);
 
     static void clean(const std::shared_ptr<ov::Node>& node);

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");
@@ -214,6 +218,11 @@ VectorDims get_planar_vdims(const snippets::lowered::ExpressionPort& expr_port);
  * @return preordered shape: `shape[i]` = `planar_shape[order[i]]` where `shape` is shape before applying the order.
  */
 VectorDims get_preordered_vdims(const snippets::lowered::ExpressionPort& expr_port);
+/**
+ * @brief Returns subtensor projected on current shape: FULL_DIM subtensor values are replaced with actual shape value
+ * @param expr_port Port whose subtensor should be processed
+ */
+VectorDims get_projected_subtensor(const snippets::lowered::ExpressionPort& expr_port);
 /* --------------------------- */
 
 /**