[Snippets][CPU][TESTS] Refactor BufferAllocation test to reduce code duplication (#31220)

### Tickets:
related to 74841

Author: aobolensk

src/common/snippets/tests/CMakeLists.txt

@@ -30,8 +30,10 @@ ov_build_target_faster(${TARGET_NAME}
 add_library(snippets_test_utils STATIC ${CMAKE_CURRENT_SOURCE_DIR}/include/lowering_utils.hpp
                                        ${CMAKE_CURRENT_SOURCE_DIR}/include/lir_test_utils.hpp
                                        ${CMAKE_CURRENT_SOURCE_DIR}/include/lir_comparator.hpp
+                                       ${CMAKE_CURRENT_SOURCE_DIR}/include/lowered/pass/buffer_allocation.hpp
                                        ${CMAKE_CURRENT_SOURCE_DIR}/src/lowering_utils.cpp
                                        ${CMAKE_CURRENT_SOURCE_DIR}/src/lir_test_utils.cpp
-                                       ${CMAKE_CURRENT_SOURCE_DIR}/src/lir_comparator.cpp)
+                                       ${CMAKE_CURRENT_SOURCE_DIR}/src/lir_comparator.cpp
+                                       ${CMAKE_CURRENT_SOURCE_DIR}/src/lowered/pass/buffer_allocation.cpp)
 target_include_directories(snippets_test_utils PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
 target_link_libraries(snippets_test_utils PRIVATE common_test_utils ov_snippets_models)

src/common/snippets/tests/include/lowered/pass/buffer_allocation.hpp

@@ -8,17 +8,19 @@
 
 #include "snippets/op/brgemm.hpp"
 #include "snippets/lowered/pass/pass.hpp"
+#include "snippets/shape_inference/shape_inference.hpp"
 
 namespace ov {
 namespace test {
 namespace snippets {
 
 typedef std::tuple<
-    bool,   // Optimized pipeline
-    bool,   // With SplitLoops opt
-    size_t, // Expected Buffer size in bytes
-    size_t, // Expected unique Buffer reg group count
-    size_t  // Expected unique Buffer cluster count
+    std::vector<ov::PartialShape>, // Input shapes
+    bool,                          // Optimized pipeline
+    bool,                          // With SplitLoops opt
+    size_t,                        // Expected Buffer size in bytes
+    size_t,                        // Expected unique Buffer reg group count
+    size_t                         // Expected unique Buffer cluster count
 > BufferAllocationParams;
 
 class BufferAllocationTest : public testing::TestWithParam<BufferAllocationParams> {
@@ -28,16 +30,23 @@ class BufferAllocationTest : public testing::TestWithParam<BufferAllocationParam
 
 protected:
     void SetUp() override;
-    void ApplyTransformations(const std::shared_ptr<ov::snippets::lowered::pass::PassConfig>& pass_config);
+    virtual void ApplyTransformations(const std::shared_ptr<ov::snippets::lowered::pass::PassConfig>& pass_config);
     void Validate();
 
-    virtual std::shared_ptr<ov::Model> GetModel() const = 0;
+    virtual std::shared_ptr<ov::Model> GetModel(const std::vector<ov::PartialShape>& shapes) const = 0;
     virtual std::shared_ptr<ov::snippets::lowered::pass::PassConfig> GetPassConfig();
+    virtual std::shared_ptr<ov::snippets::IShapeInferSnippetsFactory> GetShapeInferFactory() const;
+    virtual std::vector<ov::snippets::lowered::pass::PassPipeline::PositionedPassLowered> getBackendSpecificPasses();
 
     static void MarkOp(const std::shared_ptr<ov::Node>& node, const std::vector<size_t>& subtensor);
+    static void MarkOp(const std::shared_ptr<ov::Node>& node,
+                       const std::vector<std::vector<size_t>>& in_subtensors,
+                       const std::vector<std::vector<size_t>>& out_subtensors);
 
     ov::snippets::lowered::LinearIR m_linear_ir;
 
+    std::vector<ov::PartialShape> m_shapes;
+
     size_t m_expected_size = 0;
     size_t m_expected_reg_group_count = 0;
     size_t m_expected_cluster_count = 0;
@@ -51,7 +60,7 @@ class BufferAllocationTest : public testing::TestWithParam<BufferAllocationParam
 
 class EltwiseBufferAllocationTest : public BufferAllocationTest {
 protected:
-    std::shared_ptr<ov::Model> GetModel() const override;
+    std::shared_ptr<ov::Model> GetModel(const std::vector<ov::PartialShape>& shapes) const override;
 };
 
 }  // namespace snippets

src/common/snippets/tests/src/lowered/pass/buffer_allocation.cpp

@@ -6,7 +6,9 @@
 
 #include "openvino/opsets/opset1.hpp"
 #include "snippets/op/buffer.hpp"
+#include "snippets/op/subgraph.hpp"
 #include "snippets/lowered/linear_ir.hpp"
+#include "snippets/pass/positioned_pass.hpp"
 #include "snippets/lowered/pass/mark_loops.hpp"
 #include "snippets/lowered/pass/init_loops.hpp"
 #include "snippets/lowered/pass/insert_load_store.hpp"
@@ -25,12 +27,14 @@ namespace test {
 namespace snippets {
 
 std::string BufferAllocationTest::getTestCaseName(testing::TestParamInfo<ov::test::snippets::BufferAllocationParams> obj) {
+    std::vector<ov::PartialShape> shapes;
     bool is_optimized, with_split_loops;
     size_t expected_size, expected_reg_group_count, expected_cluster_count;
 
-    std::tie(is_optimized, with_split_loops, expected_size, expected_reg_group_count, expected_cluster_count) = obj.param;
+    std::tie(shapes, is_optimized, with_split_loops, expected_size, expected_reg_group_count, expected_cluster_count) = obj.param;
 
     std::ostringstream result;
+    result << "Shapes=" << ov::test::utils::partialShape2str(shapes) << "_";
     result << "Opt=" << ov::test::utils::bool2str(is_optimized) << "_";
     result << "Split=" << ov::test::utils::bool2str(with_split_loops) << "_";
     result << "ExpBufferSize=" << expected_size << "_";
@@ -40,13 +44,12 @@ std::string BufferAllocationTest::getTestCaseName(testing::TestParamInfo<ov::tes
 }
 
 void BufferAllocationTest::SetUp() {
-    std::tie(m_is_buffer_optimized, m_with_split_loops, m_expected_size, m_expected_reg_group_count, m_expected_cluster_count) = this->GetParam();
+    std::tie(m_shapes, m_is_buffer_optimized, m_with_split_loops, m_expected_size,
+             m_expected_reg_group_count, m_expected_cluster_count) = this->GetParam();
 
-    const auto body = GetModel();
-    m_linear_ir = ov::snippets::lowered::LinearIR(body, std::make_shared<ov::snippets::IShapeInferSnippetsFactory>());
+    const auto body = GetModel(m_shapes);
+    m_linear_ir = ov::snippets::lowered::LinearIR(body, GetShapeInferFactory());
     m_linear_ir.set_loop_depth(m_loop_depth);
-    // When Subgraph::control_flow_transformations become public method,
-    // please use this method instead of ApplyTransformations
     ApplyTransformations(GetPassConfig());
 }
 
@@ -57,6 +60,10 @@ std::shared_ptr<ov::snippets::lowered::pass::PassConfig> BufferAllocationTest::G
     return config;
 }
 
+std::shared_ptr<ov::snippets::IShapeInferSnippetsFactory> BufferAllocationTest::GetShapeInferFactory() const {
+    return std::make_shared<ov::snippets::IShapeInferSnippetsFactory>();
+}
+
 void BufferAllocationTest::MarkOp(const std::shared_ptr<ov::Node>& node, const std::vector<size_t>& subtensor) {
     for (const auto& input : node->inputs())
         ov::snippets::lowered::PortDescriptorUtils::set_port_descriptor_ptr(
@@ -66,6 +73,23 @@ void BufferAllocationTest::MarkOp(const std::shared_ptr<ov::Node>& node, const s
             output, std::make_shared<ov::snippets::lowered::PortDescriptor>(output, subtensor));
 }
 
+void BufferAllocationTest::MarkOp(const std::shared_ptr<ov::Node>& node,
+                                  const std::vector<std::vector<size_t>>& in_subtensors,
+                                  const std::vector<std::vector<size_t>>& out_subtensors) {
+    OPENVINO_ASSERT(in_subtensors.size() == node->inputs().size(), "Incorrect count of input subtensors");
+    OPENVINO_ASSERT(out_subtensors.size() == node->outputs().size(), "Incorrect count of output subtensors");
+    for (size_t i = 0; i < node->inputs().size(); ++i) {
+        const auto& input = node->input(i);
+        ov::snippets::lowered::PortDescriptorUtils::set_port_descriptor_ptr(
+            input, std::make_shared<ov::snippets::lowered::PortDescriptor>(input, in_subtensors[i]));
+    }
+    for (size_t i = 0; i < node->outputs().size(); ++i) {
+        const auto& output = node->output(i);
+        ov::snippets::lowered::PortDescriptorUtils::set_port_descriptor_ptr(
+            output, std::make_shared<ov::snippets::lowered::PortDescriptor>(output, out_subtensors[i]));
+    }
+}
+
 void BufferAllocationTest::ApplyTransformations(const std::shared_ptr<ov::snippets::lowered::pass::PassConfig>& pass_config) {
     ov::snippets::lowered::pass::PassPipeline pipeline(pass_config);
     pipeline.register_pass<ov::snippets::lowered::pass::MarkLoops>(m_vector_size);
@@ -77,8 +101,13 @@ void BufferAllocationTest::ApplyTransformations(const std::shared_ptr<ov::snippe
     pipeline.register_pass<ov::snippets::lowered::pass::InitLoops>();
     pipeline.register_pass<ov::snippets::lowered::pass::InsertLoops>();
     pipeline.register_pass<ov::snippets::lowered::pass::AllocateBuffers>(m_is_buffer_optimized);
+    pipeline.register_positioned_passes(getBackendSpecificPasses());
     pipeline.run(m_linear_ir);
 }
+std::vector<ov::snippets::lowered::pass::PassPipeline::PositionedPassLowered>
+BufferAllocationTest::getBackendSpecificPasses() {
+    return {};
+}
 
 void BufferAllocationTest::Validate() {
     std::set<size_t> reg_groups, clusters;
@@ -91,12 +120,13 @@ void BufferAllocationTest::Validate() {
     EXPECT_EQ(m_linear_ir.get_static_buffer_scratchpad_size(), m_expected_size);
 }
 
-std::shared_ptr<ov::Model> EltwiseBufferAllocationTest::GetModel() const {
+std::shared_ptr<ov::Model> EltwiseBufferAllocationTest::GetModel(const std::vector<ov::PartialShape>& shapes) const {
     const auto subtensor_eltwise = std::vector<size_t>{1, m_vector_size};
     const auto subtensor_buffer = std::vector<size_t>(2, ov::snippets::utils::get_full_dim_value());
 
-    const auto parameter0 = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape({1, 3, 100, 100}));
-    const auto parameter1 = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape({1, 3, 100, 100}));
+    const auto parameter0 = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, shapes[0]);
+    const auto parameter1 = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, shapes[0]);
+
     const auto add = std::make_shared<ov::op::v1::Add>(parameter0, parameter1);
     const auto buffer0 = std::make_shared<ov::snippets::op::Buffer>(add);
     const auto relu = std::make_shared<ov::op::v0::Relu>(buffer0);
@@ -121,6 +151,7 @@ namespace BufferAllocationTest_Instances {
 
 INSTANTIATE_TEST_SUITE_P(smoke_Snippets_BufferAllocation_EltwiseNotOptimized, EltwiseBufferAllocationTest,
                          ::testing::Combine(
+                                 ::testing::Values(std::vector<ov::PartialShape>{{1, 3, 100, 100}}),
                                  ::testing::Values(false),
                                  ::testing::Values(false),  // in this test it doesn't make sense
                                  ::testing::Values(80000),  // Each Buffer has own allocated memory
@@ -130,6 +161,7 @@ INSTANTIATE_TEST_SUITE_P(smoke_Snippets_BufferAllocation_EltwiseNotOptimized, El
 
 INSTANTIATE_TEST_SUITE_P(smoke_Snippets_BufferAllocation_EltwiseOptimized, EltwiseBufferAllocationTest,
                          ::testing::Combine(
+                                 ::testing::Values(std::vector<ov::PartialShape>{{1, 3, 100, 100}}),
                                  ::testing::Values(true),
                                  ::testing::Values(false),  // in this test it doesn't make sense
                                  ::testing::Values(40000),  // Two Buffer reuse memory