[ie transformations] improve SoftMax fusion for better mixed precision inference (#16574)

* improve SoftMax fusion

* style and unit-test fix

* more precise SoftMax unit-tests

* rewritten SoftMaxFusion with single matcher

* fixes for align_mixed_fp32_fp16_types_test.cpp and mark_subgraph_to_keep_in_mixed_precision_test.cpp

* add include for pass/pattern/op/or.hpp

* get rank only when necessary

* style-fix

* add comment why SoftmaxFusion is called manually

* fix copy_runtime_info

src/common/transformations/include/transformations/common_optimizations/softmax_fusion.hpp

@@ -17,7 +17,7 @@ class TRANSFORMATIONS_API SoftmaxFusion;
 
 /**
  * @ingroup ie_transformation_common_api
- * @brief SoftmaxFusion transformation replaces following graph:
+ * @brief SoftmaxFusion transformation replaces following graphs:
  *
  *            +---------------+
  *            │               │
@@ -63,6 +63,38 @@ class TRANSFORMATIONS_API SoftmaxFusion;
  *             │             │
  *             +-------------+
  *
+ *  and
+ *            +---------------+
+ *            │               │
+ *            │     input     │
+ *            │               │
+ *            +---------------+
+ *                    |
+ *                    |
+ *                    |
+ *                    v
+ *            +---------------+
+ *            │               │
+ *            │      Exp      │
+ *            │               │
+ *            +---------------+
+ *                │      │
+ *                │      v
+ *                │ +-----------+
+ *                │ │           │
+ *                │ │ ReduceSum │
+ *                │ │           │
+ *                │ +-----------+
+ *                │      │
+ *                │      │
+ *                v      v
+ *             +-------------+
+ *             |             │
+ *             |     Div     │
+ *             │             │
+ *             +-------------+
+ *
+ *
  * to a single Softmax node
  *
  * * Restrictions:

src/common/transformations/src/transformations/common_optimizations/mark_subgraphs_to_keep_in_mixed_precision.cpp

@@ -337,18 +337,18 @@ bool MarkSugraphsToKeepInMixedPrecision::run_on_model(const shared_ptr<ov::Model
     Manager manager(get_pass_config());
     // Mark root of Division with eps pattern to keep in FP32
     REGISTER_PASS(manager, MarkDivWithEps)
-
     REGISTER_PASS(manager, MarkExpInReduceOpPath)
-    REGISTER_PASS(manager, MarkNormalizationOps)
 
     // both Up and Down propagations are needed.
     // Why both of them are needed is explained in comments in passes declarations.
     REGISTER_PASS(manager, PropagateDownMarkToKeepInMixedPrecision)
+
     auto propagate_up = manager.register_pass<BackwardGraphRewrite>();
     ADD_MATCHER(propagate_up, PropagateUpMarkToKeepInMixedPrecision)
 
     // Mark nodes in ShapeOf subgraphs to keep in FP32
     REGISTER_PASS(manager, MarkPrecisionSensitiveShapeOfSubgraphs)
+    REGISTER_PASS(manager, MarkNormalizationOps)
     manager.run_passes(m);
 
     for (auto& node : m->get_ops()) {

src/common/transformations/src/transformations/common_optimizations/softmax_fusion.cpp

@@ -8,6 +8,7 @@
 #include <ngraph/pattern/op/wrap_type.hpp>
 #include <ngraph/rt_info.hpp>
 #include <openvino/opsets/opset6.hpp>
+#include <openvino/pass/pattern/op/or.hpp>
 #include <vector>
 
 #include "itt.hpp"
@@ -15,56 +16,66 @@
 
 ov::pass::SoftmaxFusion::SoftmaxFusion() {
     MATCHER_SCOPE(SoftmaxFusion);
-    auto data_pattern = pass::pattern::any_input(pattern::has_static_rank());
+
+    auto data_pattern = pass::pattern::any_input(pass::pattern::has_static_rank());
     auto reduce_max_axes_pattern = ngraph::pattern::wrap_type<opset6::Constant>();
     auto reduce_max_pattern = ngraph::pattern::wrap_type<opset6::ReduceMax>({data_pattern, reduce_max_axes_pattern});
     auto sub_pattern = ngraph::pattern::wrap_type<opset6::Subtract>({data_pattern, reduce_max_pattern});
-    auto exp_pattern = ngraph::pattern::wrap_type<opset6::Exp>({sub_pattern});
+
+    auto exp_input = std::make_shared<pattern::op::Or>(OutputVector{sub_pattern, data_pattern});
+    auto exp_pattern = ngraph::pattern::wrap_type<opset6::Exp>({exp_input});
+
     auto reduce_sum_axes_pattern = ngraph::pattern::wrap_type<opset6::Constant>();
     auto reduce_sum_pattern = ngraph::pattern::wrap_type<opset6::ReduceSum>({exp_pattern, reduce_sum_axes_pattern});
     auto div_pattern = ngraph::pattern::wrap_type<opset6::Divide>({exp_pattern, reduce_sum_pattern});
 
-    ov::matcher_pass_callback callback = [=](pattern::Matcher& m) {
+    ov::matcher_pass_callback callback = [=](pass::pattern::Matcher& m) {
         if (transformation_callback(m.get_match_root()))
             return false;
 
         const auto& pattern_map = m.get_pattern_value_map();
 
-        auto reduce_max_axes =
-            std::dynamic_pointer_cast<opset6::Constant>(pattern_map.at(reduce_max_axes_pattern).get_node_shared_ptr());
-        if (!reduce_max_axes || shape_size(reduce_max_axes->get_shape()) != 1)
-            return false;
         auto reduce_sum_axes =
             std::dynamic_pointer_cast<opset6::Constant>(pattern_map.at(reduce_sum_axes_pattern).get_node_shared_ptr());
         if (!reduce_sum_axes || shape_size(reduce_sum_axes->get_shape()) != 1)
             return false;
-
-        int64_t reduce_max_axis = reduce_max_axes->cast_vector<int64_t>()[0];
         int64_t reduce_sum_axis = reduce_sum_axes->cast_vector<int64_t>()[0];
-        if (reduce_max_axis < 0 || reduce_sum_axis < 0) {
-            const auto& pshape = pattern_map.at(data_pattern).get_partial_shape();
-            auto rank = pshape.rank().get_length();
-            if (reduce_max_axis < 0)
-                reduce_max_axis += rank;
-            if (reduce_sum_axis < 0)
-                reduce_sum_axis += rank;
+        if (reduce_sum_axis < 0) {
+            const auto& rank = pattern_map.at(data_pattern).get_partial_shape().rank().get_length();
+            reduce_sum_axis += rank;
+        }
+
+        auto exp_input_is_subtract = pattern_map.count(sub_pattern) != 0;
+        if (exp_input_is_subtract) {
+            auto reduce_max_axes = std::dynamic_pointer_cast<opset6::Constant>(
+                pattern_map.at(reduce_max_axes_pattern).get_node_shared_ptr());
+            if (!reduce_max_axes || shape_size(reduce_max_axes->get_shape()) != 1)
+                return false;
+            int64_t reduce_max_axis = reduce_max_axes->cast_vector<int64_t>()[0];
+
+            if (reduce_max_axis < 0) {
+                const auto& rank = pattern_map.at(data_pattern).get_partial_shape().rank().get_length();
+                reduce_max_axis += rank;
+            }
+
+            if (reduce_max_axis != reduce_sum_axis)
+                return false;
         }
-        if (reduce_max_axis != reduce_sum_axis)
-            return false;
 
         auto softmax = register_new_node<opset6::Softmax>(pattern_map.at(data_pattern), reduce_sum_axis);
         auto div = pattern_map.at(div_pattern).get_node_shared_ptr();
         softmax->set_friendly_name(div->get_friendly_name());
 
-        copy_runtime_info(
-            {
-                pattern_map.at(reduce_max_pattern).get_node_shared_ptr(),
-                pattern_map.at(sub_pattern).get_node_shared_ptr(),
-                pattern_map.at(exp_pattern).get_node_shared_ptr(),
-                pattern_map.at(reduce_sum_pattern).get_node_shared_ptr(),
-                div,
-            },
-            softmax);
+        NodeVector fused_nodes;
+        fused_nodes.push_back(pattern_map.at(exp_pattern).get_node_shared_ptr());
+        fused_nodes.push_back(pattern_map.at(reduce_sum_pattern).get_node_shared_ptr());
+        fused_nodes.push_back(div);
+        if (exp_input_is_subtract) {
+            fused_nodes.push_back(pattern_map.at(reduce_max_pattern).get_node_shared_ptr());
+            fused_nodes.push_back(pattern_map.at(sub_pattern).get_node_shared_ptr());
+        }
+        copy_runtime_info(fused_nodes, softmax);
+
         replace_node(div, softmax);
 
         return true;

src/common/transformations/tests/common_optimizations/align_mixed_fp32_fp16_types_test.cpp

@@ -171,10 +171,10 @@ TEST_F(TransformationTestsF, align_mixed_fp16_fp32_4) {
         auto reduction_axes = Constant::create(element::i64, Shape{1}, {-1});
         auto convert_to_f32_1 = make_shared<Convert>(input_1, element::f32);
         auto mvn_1 = make_shared<MVN>(convert_to_f32_1, reduction_axes, true, 1.0e-8f, op::MVNEpsMode::INSIDE_SQRT);
+        auto convert_to_f16_1 = make_shared<Convert>(mvn_1, element::f32);
         auto addition_const = Constant::create(element::f32, Shape{1}, {0.1f});
-        auto add_1 = make_shared<Add>(mvn_1, addition_const);
-        auto convert_to_f16_1 = make_shared<Convert>(add_1, element::f32);
-        auto matmul_1 = make_shared<MatMul>(convert_to_f16_1, input_2);
+        auto add_1 = make_shared<Add>(convert_to_f16_1, addition_const);
+        auto matmul_1 = make_shared<MatMul>(add_1, input_2);
 
         model_ref = make_shared<Model>(NodeVector{matmul_1}, ParameterVector{input_1, input_2});
     }
@@ -209,17 +209,16 @@ TEST_F(TransformationTestsF, align_mixed_fp16_fp32_mnv_with_split) {
         auto input_1 = make_shared<Parameter>(element::f32, Shape{1, 3, 224, 224});
         auto input_2 = make_shared<Parameter>(element::f32, Shape{1, 3, 56, 224});
 
-        auto convert_to_f32_1 = make_shared<Convert>(input_1, element::f32);
-
         auto split_axis = Constant::create(element::i64, Shape{}, {3});
-        auto split = make_shared<Split>(convert_to_f32_1, split_axis, 4);
+        auto split = make_shared<Split>(input_1, split_axis, 4);
 
         auto reduction_axes = Constant::create(element::i64, Shape{1}, {-1});
-        auto mvn_1 = make_shared<MVN>(split->output(0), reduction_axes, true, 1.0e-8f, op::MVNEpsMode::INSIDE_SQRT);
+        auto convert_to_f32_1 = make_shared<Convert>(split->output(0), element::f32);
+        auto mvn_1 = make_shared<MVN>(convert_to_f32_1, reduction_axes, true, 1.0e-8f, op::MVNEpsMode::INSIDE_SQRT);
+        auto convert_to_f16_1 = make_shared<Convert>(mvn_1, element::f32);
         auto addition_const = Constant::create(element::f32, Shape{1}, {0.1f});
-        auto add_1 = make_shared<Add>(mvn_1, addition_const);
-        auto convert_to_f16_1 = make_shared<Convert>(add_1, element::f32);
-        auto matmul_1 = make_shared<MatMul>(convert_to_f16_1, input_2);
+        auto add_1 = make_shared<Add>(convert_to_f16_1, addition_const);
+        auto matmul_1 = make_shared<MatMul>(add_1, input_2);
 
         // todo: without Converts to fp16 because of GPU
         auto result_1 = make_shared<Result>(matmul_1);

src/common/transformations/tests/common_optimizations/mark_subgraph_to_keep_in_mixed_precision_test.cpp

@@ -442,8 +442,6 @@ TEST(TransformationTests, keep_precission_sensitive_fp32_4) {
         auto matmul_1 = make_shared<MatMul>(unsqueeze_1, input_2);
 
         // marking nodes to be kept in fp32 for mixed precision
-        disable_fp16_compression(addition_const);
-        disable_fp16_compression(unsqueeze_1);
         disable_fp16_compression(mvn_1);
 
         model_ref = make_shared<Model>(NodeVector{matmul_1}, ParameterVector{input_1, input_2});
@@ -487,8 +485,6 @@ TEST(TransformationTests, keep_precission_sensitive_fp32_5) {
         auto matmul_1 = make_shared<MatMul>(unsqueeze_1, input_2);
 
         // marking nodes to be kept in fp32 for mixed precision
-        disable_fp16_compression(addition_const);
-        disable_fp16_compression(unsqueeze_1);
         disable_fp16_compression(normalizel2_1);
 
         model_ref = make_shared<Model>(NodeVector{matmul_1}, ParameterVector{input_1, input_2});

src/common/transformations/tests/common_optimizations/softmax_fusion.cpp

@@ -17,6 +17,7 @@ using namespace testing;
 using namespace ngraph;
 
 class SoftmaxFusionFixture : public ::testing::TestWithParam<std::tuple<int64_t, int64_t>> {};
+class SoftmaxFusionSimplePatternFixture : public ::testing::TestWithParam<std::tuple<int64_t>> {};
 
 TEST_P(SoftmaxFusionFixture, SoftmaxFusion) {
     Shape shape{1, 1, 256};
@@ -67,25 +68,16 @@ INSTANTIATE_TEST_SUITE_P(SoftmaxFusionTests,
                                            std::make_tuple(2, -1),
                                            std::make_tuple(2, 2)));
 
-class NegativeSoftmaxFusionFixture
-    : public ::testing::TestWithParam<std::tuple<std::vector<int64_t>, std::vector<int64_t>>> {};
-
-TEST_P(NegativeSoftmaxFusionFixture, NegativeSoftmaxFusion) {
-    Shape shape{1, 1, 256};
+TEST_P(SoftmaxFusionSimplePatternFixture, SoftmaxFusionSimplePatternTest) {
+    Shape shape{1, 3, 256, 256};
     auto params = GetParam();
-    auto reduce_max_axes_val = std::get<0>(params);
-    auto reduce_sum_axes_val = std::get<1>(params);
+    auto reduce_axis_val = std::get<0>(params);
     std::shared_ptr<Function> f(nullptr), f_ref(nullptr);
     {
         auto data = std::make_shared<opset6::Parameter>(element::f32, shape);
-        auto reduce_max_axes =
-            opset6::Constant::create(element::i64, Shape{reduce_max_axes_val.size()}, reduce_max_axes_val);
-        auto reduce_max = std::make_shared<opset6::ReduceMax>(data, reduce_max_axes);
-        auto sub = std::make_shared<opset6::Subtract>(data, reduce_max);
-        auto exp = std::make_shared<opset6::Exp>(sub);
-        auto reduce_sum_axes =
-            opset6::Constant::create(element::i64, Shape{reduce_sum_axes_val.size()}, reduce_sum_axes_val);
-        auto reduce_sum = std::make_shared<opset6::ReduceSum>(exp, reduce_sum_axes);
+        auto exp = std::make_shared<opset6::Exp>(data);
+        auto reduce_axis = opset6::Constant::create(element::i64, Shape{}, {reduce_axis_val});
+        auto reduce_sum = std::make_shared<opset6::ReduceSum>(exp, reduce_axis, true);
         auto div = std::make_shared<opset6::Divide>(exp, reduce_sum);
         f = std::make_shared<Function>(NodeVector{div}, ParameterVector{data});
 
@@ -100,16 +92,10 @@ TEST_P(NegativeSoftmaxFusionFixture, NegativeSoftmaxFusion) {
     }
     {
         auto data = std::make_shared<opset6::Parameter>(element::f32, shape);
-        auto reduce_max_axes =
-            opset6::Constant::create(element::i64, Shape{reduce_max_axes_val.size()}, reduce_max_axes_val);
-        auto reduce_max = std::make_shared<opset6::ReduceMax>(data, reduce_max_axes);
-        auto sub = std::make_shared<opset6::Subtract>(data, reduce_max);
-        auto exp = std::make_shared<opset6::Exp>(sub);
-        auto reduce_sum_axes =
-            opset6::Constant::create(element::i64, Shape{reduce_sum_axes_val.size()}, reduce_sum_axes_val);
-        auto reduce_sum = std::make_shared<opset6::ReduceSum>(exp, reduce_sum_axes);
-        auto div = std::make_shared<opset6::Divide>(exp, reduce_sum);
-        f_ref = std::make_shared<Function>(NodeVector{div}, ParameterVector{data});
+        if (reduce_axis_val < 0)
+            reduce_axis_val += shape.size();
+        auto softmax = std::make_shared<opset6::Softmax>(data, reduce_axis_val);
+        f_ref = std::make_shared<Function>(NodeVector{softmax}, ParameterVector{data});
     }
 
     auto fc =
@@ -118,6 +104,50 @@ TEST_P(NegativeSoftmaxFusionFixture, NegativeSoftmaxFusion) {
     ASSERT_TRUE(res.valid) << res.message;
 }
 
+INSTANTIATE_TEST_SUITE_P(SoftmaxFusionSimplePatternTests,
+                         SoftmaxFusionSimplePatternFixture,
+                         ::testing::Values(std::make_tuple(0),
+                                           std::make_tuple(1),
+                                           std::make_tuple(2),
+                                           std::make_tuple(-1),
+                                           std::make_tuple(-2)));
+
+class NegativeSoftmaxFusionFixture
+    : public ::testing::TestWithParam<std::tuple<std::vector<int64_t>, std::vector<int64_t>>> {};
+
+TEST_P(NegativeSoftmaxFusionFixture, NegativeSoftmaxFusion) {
+    // ReduceMax arguments do not match conditions, therefore these nodes
+    // are not included into final SoftMax node
+    Shape shape{1, 1, 256};
+    auto params = GetParam();
+    auto reduce_max_axes_val = std::get<0>(params);
+    auto reduce_sum_axes_val = std::get<1>(params);
+    std::shared_ptr<Function> f(nullptr);
+
+    auto data = std::make_shared<opset6::Parameter>(element::f32, shape);
+    auto reduce_max_axes =
+        opset6::Constant::create(element::i64, Shape{reduce_max_axes_val.size()}, reduce_max_axes_val);
+    auto reduce_max = std::make_shared<opset6::ReduceMax>(data, reduce_max_axes);
+    auto sub = std::make_shared<opset6::Subtract>(data, reduce_max);
+    auto exp = std::make_shared<opset6::Exp>(sub);
+    auto reduce_sum_axes =
+        opset6::Constant::create(element::i64, Shape{reduce_sum_axes_val.size()}, reduce_sum_axes_val);
+    auto reduce_sum = std::make_shared<opset6::ReduceSum>(exp, reduce_sum_axes);
+    auto div = std::make_shared<opset6::Divide>(exp, reduce_sum);
+    f = std::make_shared<Function>(NodeVector{div}, ParameterVector{data});
+
+    auto unh = std::make_shared<ngraph::pass::UniqueNamesHolder>();
+    pass::Manager m;
+    m.register_pass<pass::InitUniqueNames>(unh);
+    m.register_pass<ov::pass::InitNodeInfo>();
+    m.register_pass<ov::pass::SoftmaxFusion>();
+    m.register_pass<pass::CheckUniqueNames>(unh);
+    m.run_passes(f);
+    ASSERT_NO_THROW(check_rt_info(f));
+    ASSERT_EQ(count_ops_of_type<opset6::ReduceMax>(f), 1);
+    ASSERT_EQ(count_ops_of_type<opset6::Subtract>(f), 1);
+}
+
 INSTANTIATE_TEST_SUITE_P(NegativeSoftmaxFusionTests,
                          NegativeSoftmaxFusionFixture,
                          ::testing::ValuesIn(std::vector<std::tuple<std::vector<int64_t>, std::vector<int64_t>>>{

src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp

@@ -44,6 +44,7 @@
 #include "transformations/common_optimizations/convert_compression_only_to_legacy.hpp"
 #include <transformations/common_optimizations/wrap_interpolate_into_transposes.hpp>
 #include <transformations/common_optimizations/transpose_sinking.hpp>
+#include <transformations/common_optimizations/softmax_fusion.hpp>
 
 #include <transformations/op_conversions/convert_depth_to_space.hpp>
 #include <transformations/op_conversions/convert_space_to_depth.hpp>
@@ -190,6 +191,10 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
 
         type_to_fuse_map empty_fuse_map = {};
         manager.register_pass<ov::pass::Validate>();
+
+        // fuse softmax patterns so that they will not be marked as precision sensitive in ConvertPrecision
+        manager.register_pass<ov::pass::SoftmaxFusion>();
+
         //  call ConvertPrecision with keep_precision_sensitive_in_fp32 = true
         manager.register_pass<ov::pass::ConvertPrecision>(fp_convert_precision_map, empty_fuse_map, true);