Extended tests coverage

v-Golubev · Aug 2, 2024 · dbf9865 · dbf9865
1 parent 613d0c9
commit dbf9865
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 8 deletions.
diff --git a/src/plugins/intel_cpu/src/emitters/snippets/cpu_runtime_configurator.cpp b/src/plugins/intel_cpu/src/emitters/snippets/cpu_runtime_configurator.cpp
@@ -164,14 +164,38 @@ std::unordered_set<UnifiedLoopInfoPtr> CPURuntimeConfigurator::ParallelWAOptimiz
 }
 
 bool CPURuntimeConfigurator::ParallelWAOptimizer::check_brgemms(const ov::snippets::lowered::LinearIRPtr& linear_ir) {
-    auto found_it = std::find_if(linear_ir->begin(), linear_ir->end(), [](const ExpressionPtr& expr) {
+    auto is_brgemm = [](const ExpressionPtr& expr) {
         return ov::is_type<ov::snippets::op::Brgemm>(expr->get_node());
-    });
-    if (found_it == linear_ir->end())
+    };
+    auto brgemm_it = std::find_if(linear_ir->begin(), linear_ir->end(), is_brgemm);
+    std::unordered_set<ExpressionPtr> brgemms;
+    while (brgemm_it != linear_ir->end()) {
+        brgemms.insert(*brgemm_it);
+        brgemm_it = std::find_if(std::next(brgemm_it), linear_ir->end(), is_brgemm);
+    }
+    if (brgemms.empty())
         return false;
-    const auto& brgemm = *found_it;
-    const auto planar_shape = ov::snippets::utils::get_planar_vdims(brgemm->get_input_port(0));
-    return ov::snippets::utils::is_dynamic_value(*++planar_shape.rbegin());
+
+    const auto loop_manager = linear_ir->get_loop_manager();
+    for (const auto& brgemm : brgemms) {
+        const auto& loop_idces = brgemm->get_loop_ids();
+        if (loop_idces.empty())
+            return false;
+        const auto& outermost_loop_id = loop_idces[0];
+        const auto& outermost_loop = loop_manager->get_loop_info(outermost_loop_id);
+        // WA: the feature doesn't work in static case
+        if (!ov::snippets::utils::is_dynamic_value(outermost_loop->get_work_amount())) {
+            return false;
+        }
+        bool loop_by_m = true;
+        outermost_loop->iterate_through_ports([&loop_by_m](const LoopPort& port) {
+            if (port.is_incremented && port.dim_idx != 1)
+                loop_by_m = false;
+        });
+        if (!loop_by_m)
+            return false;
+    }
+    return true;
 }
 
 void CPURuntimeConfigurator::ParallelWAOptimizer::init(const ov::snippets::lowered::LinearIRPtr& linear_ir) {

diff --git a/src/plugins/intel_cpu/tests/functional/shared_tests_instances/snippets/mha.cpp b/src/plugins/intel_cpu/tests/functional/shared_tests_instances/snippets/mha.cpp
@@ -142,9 +142,21 @@ std::vector<std::vector<ov::test::InputShape>> splitm_dynamic_shapes_4d = {
         {PartialShape{-1, -1, -1, -1}, {{1, 128, 2, 64}, {1, 17, 2, 64}, {1, 128, 2, 64}}},
     },
     {
-        {PartialShape{-1, -1, -1, -1}, {{1, 32, 2, 64}}},
+        {PartialShape{-1, 128, -1, -1}, {{1, 128, 2, 64}}},
         {PartialShape{-1, -1, -1, -1}, {{1, 16, 2, 64}}},
-        {PartialShape{-1, -1, -1, -1}, {{1, 1, 32, 16}}},
+        {PartialShape{-1, -1, 128, -1}, {{1, 1, 128, 16}}},
+        {PartialShape{-1, -1, -1, -1}, {{1, 16, 2, 32}}},
+    },
+    {
+        {PartialShape{-1, 32, -1, -1}, {{1, 32, 2, 64}}},
+        {PartialShape{-1, -1, -1, -1}, {{1, 16, 2, 64}}},
+        {PartialShape{-1, -1, 32, -1}, {{1, 1, 32, 16}}},
+        {PartialShape{-1, -1, -1, -1}, {{1, 16, 2, 32}}},
+    },
+    {
+        {PartialShape{-1, -1, -1, -1}, {{1, 16, 2, 64}}},
+        {PartialShape{-1, -1, -1, -1}, {{1, 16, 2, 64}}},
+        {PartialShape{-1, -1, 16, -1}, {{1, 1, 16, 16}}},
         {PartialShape{-1, -1, -1, -1}, {{1, 16, 2, 32}}},
     },
 };