Merge remote-tracking branch 'origin/master' into readability-avoid-nested-conditional-operator

Author: aobolensk

docs/requirements.txt

@@ -1,13 +1,13 @@
 alabaster==0.7.16
 atomicwrites==1.4.1
 attrs==22.2.0
-Babel==2.11.0
+Babel==2.17.0
 beautifulsoup4==4.9.3
 breathe==4.36.0
 certifi==2024.7.4
 colorama==0.4.6
 Cython==3.0.12
-docutils==0.20
+docutils==0.21.2
 idna==3.7
 imagesize==1.3.0
 importlib-metadata==4.8.0
@@ -18,9 +18,9 @@ jsonschema==4.23.0
 lxml>=4.9.2
 MarkupSafe==2.1.1
 mistune==2.0.3
-myst-parser==2.0.0
+myst-parser==4.0.1
 packaging==23.0
-pluggy==0.13.1
+pluggy==1.6.0
 pydata-sphinx-theme==0.14.4
 Pygments==2.15.1
 pyparsing==3.0.9

src/common/snippets/include/snippets/emitter.hpp

@@ -16,7 +16,15 @@ namespace snippets {
  * @interface RegType
  * @brief Register type of input and output operations
  */
-enum class RegType { gpr, vec, mask, undefined };
+enum class RegType {
+    gpr,
+    vec,
+    mask,
+    // Ticket: 166071
+    // Need to move this type to a separate class
+    address,  // address type should be ignored by the code generation logic, as it is handled outside the snippets pipeline.
+    undefined
+};
 /**
  * @interface Reg
  * @brief Register representation: type of register and index
@@ -26,7 +34,8 @@ struct Reg {
     Reg() = default;
     Reg(RegType type_, size_t idx_) : type(type_), idx(idx_) {}
 
-    bool is_defined() const  { return  type != RegType::undefined && idx != UNDEFINED_IDX; }
+    bool is_address() const { return type == RegType::address; }
+    bool is_defined() const { return is_address() || (type != RegType::undefined && idx != UNDEFINED_IDX); }
     RegType type = RegType::undefined;
     size_t idx = UNDEFINED_IDX;
 

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

@@ -44,9 +44,8 @@ class PortDescriptor {
     void set_shape(const VectorDims& tensor);
     void set_layout(const std::vector<size_t>& layout);
     void set_subtensor(const VectorDims& subtensor);
-    void set_reg(Reg reg) { m_reg = std::move(reg); }
-    void set_reg_type(RegType type) { m_reg.type = type; }
-    void set_reg_idx(size_t idx) { m_reg.idx = idx; }
+    void set_reg(Reg reg);
+    void set_reg_type(RegType type);
 
     // Indexing starts from the end (rbegin() + idx)
     void set_subtensor_dim(size_t idx, VectorDims::value_type value);
@@ -92,6 +91,9 @@ class PortDescriptorUtils {
     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 set_address_reg_type(const ov::Input<ov::Node>& in);
+    static void set_address_reg_type(const ov::Output<ov::Node>& out);
+
     static void clean(const std::shared_ptr<ov::Node>& node);
 
 private:

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

@@ -41,7 +41,7 @@ class Brgemm : virtual public modifier::MemoryAccess, public ov::op::Op {
     bool visit_attributes(AttributeVisitor& visitor) override;
 
 protected:
-    ov::element::Type get_output_type() const;
+    virtual ov::element::Type get_output_type() const;
     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;

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

@@ -34,6 +34,13 @@ class Scalar  : public ov::op::v0::Constant {
     }
     Scalar& operator=(const Scalar&) = delete;
 
+    template <class T, class = typename std::enable_if<std::is_fundamental<T>::value>::type>
+    T get_value() const {
+        const auto vec = cast_vector<T>();
+        OPENVINO_ASSERT(vec.size() == 1, "Scalar must have a single value");
+        return vec[0];
+    }
+
     std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
     void validate_and_infer_types() override;
     bool visit_attributes(AttributeVisitor& visitor) override;

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

@@ -67,11 +67,23 @@ class SnippetsTokenization : public ov::pass::ModelPass {
      * @ingroup snippets
      */
     struct Config {
-        Config(size_t concurrency, size_t data_ptr_gpr_count, bool split_m_dimension, bool enable_transpose_on_output,
-               bool dyn_mha_token, std::set<size_t> mha_transpose_ranks)
-            : m_concurrency(concurrency), m_data_ptr_gpr_count(data_ptr_gpr_count), m_split_m_dimension(split_m_dimension),
-              m_mha_token_enable_transpose_on_output(enable_transpose_on_output), m_is_dynamic_mha_token_enabled(dyn_mha_token),
-              m_mha_supported_transpose_ranks(std::move(mha_transpose_ranks)) {
+        using CanBeFusedAsPostOpPred = std::function<bool(const std::shared_ptr<const ov::op::v0::MatMul>&,
+                                                          const std::shared_ptr<const ov::Node>&)>;
+
+        Config(size_t concurrency,
+               size_t data_ptr_gpr_count,
+               bool split_m_dimension,
+               bool enable_transpose_on_output,
+               bool dyn_mha_token,
+               std::set<size_t> mha_transpose_ranks,
+               CanBeFusedAsPostOpPred can_be_fused_as_postop = nullptr)
+            : m_concurrency(concurrency),
+              m_data_ptr_gpr_count(data_ptr_gpr_count),
+              m_split_m_dimension(split_m_dimension),
+              m_mha_token_enable_transpose_on_output(enable_transpose_on_output),
+              m_is_dynamic_mha_token_enabled(dyn_mha_token),
+              m_mha_supported_transpose_ranks(std::move(mha_transpose_ranks)),
+              m_can_be_fused_as_postop(can_be_fused_as_postop) {
             OPENVINO_ASSERT(concurrency > 0, "Concurrency should be greater than 0");
             OPENVINO_ASSERT(data_ptr_gpr_count > 0, "data_ptr_gpr_count should be greater than 0");
         }
@@ -104,6 +116,10 @@ class SnippetsTokenization : public ov::pass::ModelPass {
             return m_mha_supported_transpose_ranks;
         }
 
+        const CanBeFusedAsPostOpPred& get_can_be_fused_as_postop() const {
+            return m_can_be_fused_as_postop;
+        }
+
     private:
         size_t m_concurrency = 0;
         // The number of gpr that can be used as data pointers for data nodes (Parameter (and non-Scalar Constants),
@@ -123,6 +139,9 @@ class SnippetsTokenization : public ov::pass::ModelPass {
         // Note that in general Snippets support Transpose of any ranks.
         // But at the moment Transpose is used only in MHA pattern where 3D and 4D tensors are supported.
         std::set<size_t> m_mha_supported_transpose_ranks = { 3, 4 };
+        // Predicate that checks if the node can be fused as MatMul post-op.
+        // It is currently used only in TokenizeMLPSeqSnippets
+        CanBeFusedAsPostOpPred m_can_be_fused_as_postop = nullptr;
     };
 
     SnippetsTokenization(const Config& config) : m_config(config) {}

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

@@ -13,7 +13,9 @@ inline static std::vector<size_t> transform_snippets_regs_to_idxs(const std::vec
     std::vector<size_t> idxs;
     idxs.reserve(regs.size());
     for (const auto& reg : regs) {
-        OPENVINO_ASSERT(expected_type == snippets::RegType::undefined || reg.type == expected_type, "Reg type mismatch during to_idxs conversion");
+        OPENVINO_ASSERT(reg.type == snippets::RegType::address ||
+                        expected_type == snippets::RegType::undefined || reg.type == expected_type,
+                        "Reg type mismatch during to_idxs conversion");
         idxs.emplace_back(reg.idx);
     }
     return idxs;

src/common/snippets/src/emitter.cpp

@@ -33,6 +33,8 @@ std::ostream& operator<<(std::ostream& s, const Reg& r) {
                 return "mask";
             case RegType::undefined:
                  return "undefined";
+            case RegType::address:
+                 return "address";
             default:
                 OPENVINO_THROW("Unexpected RegType");
         }

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

@@ -4,12 +4,11 @@
 
 #include "snippets/lowered/pass/assign_registers.hpp"
 
+#include "snippets/itt.hpp"
 #include "snippets/lowered/linear_ir.hpp"
+#include "snippets/op/kernel.hpp"
 #include "snippets/snippets_isa.hpp"
-#include "snippets/itt.hpp"
 #include "snippets/utils/utils.hpp"
-#include "snippets/op/kernel.hpp"
-
 
 // This header is needed to avoid MSVC warning "C2039: 'inserter': is not a member of 'std'"
 #include <iterator>
@@ -21,9 +20,19 @@ namespace pass {
 
 AssignRegisters::RegMap AssignRegisters::assign_regs_manually(const LinearIR& linear_ir, std::set<Reg>& gpr_pool, std::set<Reg>& vec_pool) {
     RegMap manually_assigned;
-    OPENVINO_ASSERT(gpr_pool.size() >= (linear_ir.get_parameters().size() + linear_ir.get_results().size()),
+    const auto internal_parameters = [&]() {
+        std::vector<ExpressionPtr> internal_parameters;
+        internal_parameters.reserve(linear_ir.get_parameters().size());
+        for (const auto& param : linear_ir.get_parameters()) {
+            if (!param->get_output_port_descriptor(0)->get_reg().is_address()) {
+                internal_parameters.push_back(param);
+            }
+        }
+        return internal_parameters;
+    }();
+    OPENVINO_ASSERT(gpr_pool.size() >= (internal_parameters.size() + linear_ir.get_results().size()),
                     "Not enough gp registers in the pool to perform manual assignment");
-    for (const auto& param : linear_ir.get_parameters()) {
+    for (const auto& param : internal_parameters) {
         manually_assigned[param->get_output_port_descriptor(0)->get_reg()] = *gpr_pool.begin();
         gpr_pool.erase(gpr_pool.begin());
     }
@@ -179,10 +188,16 @@ bool AssignRegisters::run(LinearIR& linear_ir) {
         for (const auto& live_reg : expr->get_live_regs())
             mapped_live_regs.insert(assigned_reg_map[live_reg]);
         expr->set_live_regs(mapped_live_regs);
-        for (const auto& in : expr->get_input_port_descriptors())
-            in->set_reg(assigned_reg_map[in->get_reg()]);
-        for (const auto& out : expr->get_output_port_descriptors())
-            out->set_reg(assigned_reg_map[out->get_reg()]);
+        for (const auto& in : expr->get_input_port_descriptors()) {
+            if (!in->get_reg().is_address()) {
+                in->set_reg(assigned_reg_map[in->get_reg()]);
+            }
+        }
+        for (const auto& out : expr->get_output_port_descriptors()) {
+            if (!out->get_reg().is_address()) {
+                out->set_reg(assigned_reg_map[out->get_reg()]);
+            }
+        }
     }
     return false;
 }

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

@@ -3,9 +3,10 @@
 //
 
 #include "snippets/lowered/pass/init_live_ranges.hpp"
+
 #include "snippets/itt.hpp"
-#include "snippets/op/subgraph.hpp"
 #include "snippets/lowered/expressions/buffer_expression.hpp"
+#include "snippets/op/subgraph.hpp"
 
 namespace ov {
 namespace snippets {
@@ -17,9 +18,9 @@ inline bool pass_through_expr(const ExpressionPtr& expr) {
     const auto& node = expr->get_node();
     return op::Subgraph::is_shape_infer_op(node)
 #ifdef SNIPPETS_DEBUG_CAPS
-            || ov::is_type_any_of<op::PerfCountBeginBase, op::PerfCountEndBase>(node)
+           || ov::is_type_any_of<op::PerfCountBeginBase, op::PerfCountEndBase>(node)
 #endif
-            || ov::is_type<BufferExpression>(expr);
+           || ov::is_type<BufferExpression>(expr);
 }
 
 } // namespace

src/common/snippets/src/lowered/port_descriptor.cpp

@@ -66,6 +66,16 @@ void PortDescriptor::set_subtensor(const VectorDims& subtensor) {
     m_subtensor_shape = subtensor;
 }
 
+void PortDescriptor::set_reg(Reg reg) {
+    OPENVINO_ASSERT(m_reg.type != RegType::address, "Failed to set reg: reg with 'address' type mustn't be changed");
+    m_reg = std::move(reg);
+}
+
+void PortDescriptor::set_reg_type(RegType type) {
+    OPENVINO_ASSERT(m_reg.type != RegType::address, "Failed to set reg type: address reg type mustn't be changed");
+    m_reg.type = type;
+}
+
 void PortDescriptor::set_subtensor_dim(size_t idx, VectorDims::value_type value) {
     OPENVINO_ASSERT(idx < m_subtensor_shape.size(), "Failed to set subtensor value: idx should be less than size");
     *(m_subtensor_shape.rbegin() + idx) = value;
@@ -202,6 +212,18 @@ PortDescriptorPtr PortDescriptorUtils::get_port_descriptor_ptr(const Output<cons
     return out_descs[out.get_index()];
 }
 
+void PortDescriptorUtils::set_address_reg_type(const ov::Input<ov::Node>& in) {
+    auto desc = get_port_descriptor_ptr(in);
+    desc->set_reg_type(RegType::address);
+    set_port_descriptor_ptr(in, desc);
+}
+
+void PortDescriptorUtils::set_address_reg_type(const ov::Output<ov::Node>& out) {
+    auto desc = get_port_descriptor_ptr(out);
+    desc->set_reg_type(RegType::address);
+    set_port_descriptor_ptr(out, desc);
+}
+
 void PortDescriptorUtils::clean(const std::shared_ptr<ov::Node>& node) {
     auto& rt_info = node->get_rt_info();
     rt_info.erase(PortDescriptorVectorAttribute::get_type_info_static());

src/common/snippets/src/pass/mlp_seq_tokenization.cpp

@@ -27,21 +27,20 @@ inline bool has_one_consumer(const std::shared_ptr<ov::Node>& node) {
 }
 
 size_t get_potential_body_params(const std::shared_ptr<ov::Node>& op) {
+    if (const auto fq = ov::as_type_ptr<ov::op::v0::FakeQuantize>(op)) {
+        return ov::snippets::utils::get_non_scalar_constant_count_for_fq(fq);
+    }
     size_t count = 0;
     for (size_t i = 1; i < op->get_input_size(); ++i) {
         const auto input = op->input_value(i);
         const auto parent = input.get_node_shared_ptr();
         const auto constant = ov::as_type_ptr<ov::op::v0::Constant>(parent);
         const auto is_scalar = constant && (ov::shape_size(input.get_shape()) == 1);
         const auto should_be_inside_body = constant && ov::snippets::op::Subgraph::constant_input_should_be_inside_body(op);
-        const auto is_fq_weight = constant && ov::is_type<ov::op::v0::FakeQuantize>(op);
-        if (!(is_scalar || should_be_inside_body || is_fq_weight)) {
+        if (!(is_scalar || should_be_inside_body)) {
             count++;
         }
     }
-    if (const auto fq = ov::as_type_ptr<ov::op::v0::FakeQuantize>(op)) {
-        count += ov::snippets::utils::get_non_scalar_constant_count_for_fq(fq);
-    }
     return count;
 }
 
@@ -150,14 +149,26 @@ TokenizeMLPSeqSnippets::TokenizeMLPSeqSnippets(const SnippetsTokenization::Confi
         auto interm_op = prev_op->get_output_target_inputs(0).begin()->get_node()->shared_from_this();
         auto available_regs = config.get_data_ptr_gpr_count();
 
+        bool postops_fusion_possible = true;
+        std::shared_ptr<ov::op::v0::MatMul> cur_matmul = matmul0;
         while (has_one_consumer(prev_op)) {
             auto current_potential_body_params_count = potential_body_params_count;
 
             if (is_matmul_supported(interm_op) && !transformation_callback(interm_op)) {
                 // +1 for weights
                 current_potential_body_params_count++;
+                // If MatMul is the first in the sequence, postops fusion status is reset
+                postops_fusion_possible = true;
+                cur_matmul = ov::as_type_ptr<ov::op::v0::MatMul>(interm_op);
+                OPENVINO_ASSERT(cur_matmul, "MatMul is expected");
             } else if (is_supported_intermediate_op(interm_op)) {
-                current_potential_body_params_count += get_potential_body_params(interm_op);
+                // Intermediate op contributes to the body params count only if can't be fused as post-op
+                // or if a previous node between MatMul and this op is not supported by post-op fusion
+                if (!postops_fusion_possible || config.get_can_be_fused_as_postop() == nullptr ||
+                    !config.get_can_be_fused_as_postop()(cur_matmul, interm_op)) {
+                    postops_fusion_possible = false;
+                    current_potential_body_params_count += get_potential_body_params(interm_op);
+                }
             } else {
                 // Unsupported op
                 break;
@@ -190,7 +201,6 @@ TokenizeMLPSeqSnippets::TokenizeMLPSeqSnippets(const SnippetsTokenization::Confi
 
         // mark the Subgraph as Completed to not allow Snippets to include any nodes into this Subgraph in common Tokenization
         SetSnippetsSubgraphType(subgraph, SnippetsSubgraphType::Completed);
-
         return true;
     });
 }

src/common/snippets/src/pass/propagate_precision.cpp

@@ -10,7 +10,6 @@
 #include "openvino/core/rt_info.hpp"
 #include "transformations/utils/utils.hpp"
 
-#include <assert.h>
 #include <memory>
 
 
@@ -65,12 +64,12 @@ bool ov::snippets::pass::PropagatePrecision::run_on_model(const std::shared_ptr<
             input_precisions.push_back(input_precision);
         }
 
-        assert(std::all_of(
-            supported_precisions.begin(),
-            supported_precisions.end(),
-            [&input_precisions](const std::vector<element::Type>& precisions) {
-                return precisions.size() == input_precisions.size();
-            }) && "input precisions count is not equal for supported precisions");
+        OPENVINO_ASSERT(std::all_of(supported_precisions.begin(),
+                                    supported_precisions.end(),
+                                    [&input_precisions](const std::vector<element::Type>& precisions) {
+                                        return precisions.size() == input_precisions.size();
+                                    }),
+                        "input precisions count is not equal for supported precisions");
 
         // update input precisions
         // if possible then convert precisions to supported

src/common/snippets/src/shape_inference/shape_infer_instances.cpp

@@ -152,7 +152,9 @@ Result HorizonOpShapeInfer::infer(const std::vector<VectorDimsRef>& input_shapes
 }
 
 BrgemmShapeInfer::BrgemmShapeInfer(const std::shared_ptr<Node>& n) {
-    for (const auto& in : n->inputs()) {
+    // Only first 2 inputs are used for shape inference
+    for (size_t i = 0; i < 2; ++i) {
+        const auto& in = n->input(i);
         const auto& port = lowered::PortDescriptorUtils::get_port_descriptor_ptr(in);
         m_io_layouts.push_back(port->get_layout());
     }
@@ -161,7 +163,7 @@ BrgemmShapeInfer::BrgemmShapeInfer(const std::shared_ptr<Node>& n) {
 }
 
 Result BrgemmShapeInfer::infer(const std::vector<VectorDimsRef>& input_shapes) {
-    OPENVINO_ASSERT(input_shapes.size() == 2 || input_shapes.size() == 3, "BRGEMM expects 2 or 3 input shapes for shape inference");
+    OPENVINO_ASSERT(input_shapes.size() >= 2, "Unexpected input_shapes count");
 
     // Todo: Ideally we should use the layout stored in PortDescriptors. Can we do it?
     const auto& arg0_shape = ov::snippets::utils::get_planar_vdims(input_shapes[0].get(), m_io_layouts[0]);