[CPU]Decompose Logsoftmax (#13679)

src/plugins/intel_cpu/src/ngraph_transformations/convert_logsoftmax.cpp

@@ -0,0 +1,39 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "convert_logsoftmax.hpp"
+#include <openvino/opsets/opset1.hpp>
+#include <openvino/opsets/opset8.hpp>
+#include <openvino/core/rt_info.hpp>
+#include <openvino/pass/pattern/op/wrap_type.hpp>
+#include <openvino/pass/pattern/op/or.hpp>
+#include <transformations/utils/utils.hpp>
+
+#include "itt.hpp"
+
+ov::intel_cpu::ConvertLogSoftmax::ConvertLogSoftmax() {
+    MATCHER_SCOPE(ConvertLogSoftmax);
+    using namespace ov::pass::pattern;
+    auto input_m = any_input();
+    auto logSoftmax_m = wrap_type<ov::opset8::LogSoftmax>({input_m});
+    matcher_pass_callback callback = [=](Matcher& m) {
+        const auto& pattern_map = m.get_pattern_value_map();
+        auto logSoftmaxNode = ov::as_type_ptr<ov::opset8::LogSoftmax>(pattern_map.at(logSoftmax_m).get_node_shared_ptr());
+        auto axis = std::make_shared<opset8::Constant>(ov::element::i64, ov::Shape{}, logSoftmaxNode->get_axis());
+        auto input = pattern_map.at(input_m);
+        auto xMax = std::make_shared<opset8::ReduceMax>(input, axis, true);
+        auto subtract = std::make_shared<opset8::Subtract>(input, xMax);
+        auto exp = std::make_shared<opset8::Exp>(subtract);
+        auto s = std::make_shared<opset8::ReduceSum>(exp, axis, true);
+        auto log = std::make_shared<opset8::Log>(s);
+        auto result = std::make_shared<opset8::Subtract>(subtract, log);
+        result->set_friendly_name(logSoftmaxNode->get_friendly_name());
+        ov::copy_runtime_info(logSoftmaxNode, {xMax, subtract, exp, s, log, result});
+        ov::replace_node(logSoftmaxNode, result);
+        return true;
+    };
+
+    auto m = std::make_shared<Matcher>(logSoftmax_m, matcher_name);
+    this->register_matcher(m, callback);
+}

src/plugins/intel_cpu/src/ngraph_transformations/convert_logsoftmax.hpp

@@ -0,0 +1,18 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <openvino/pass/graph_rewrite.hpp>
+
+namespace ov {
+namespace intel_cpu {
+class ConvertLogSoftmax: public ov::pass::MatcherPass {
+public:
+    OPENVINO_RTTI("ConvertLogSoftmax", "0");
+    ConvertLogSoftmax();
+};
+
+}   // namespace intel_cpu
+}   // namespace ov

src/plugins/intel_cpu/src/ngraph_transformations/convert_to_cpu_specific_opset.hpp

@@ -19,6 +19,7 @@
 #include "transformations/utils/utils.hpp"
 #include "rnn_sequences_optimization.hpp"
 #include "transformations/common_optimizations/reshape_sequence_fusion.hpp"
+#include "convert_logsoftmax.hpp"
 
 #include "itt.hpp"
 
@@ -29,6 +30,7 @@ inline void ConvertToCPUSpecificOpset(std::shared_ptr<ngraph::Function> &nGraphF
     RUN_ON_FUNCTION_SCOPE(ConvertToCPUSpecificOpset);
     ngraph::pass::Manager manager;
     manager.register_pass<ConvertMatMulToFC>();
+    manager.register_pass<ConvertLogSoftmax>();
     manager.register_pass<AlignMatMulInputRanks>();
     manager.register_pass<ConvertTileToSeqTiles>();
     manager.register_pass<FullyConnectedBiasFusion>();
@@ -44,6 +46,7 @@ inline void ConvertToCPUSpecificOpset(std::shared_ptr<ngraph::Function> &nGraphF
     manager.register_pass<ngraph::pass::ConstantFolding>();
     manager.register_pass<ngraph::pass::ConvertPrecision>(precisions_array {{ ngraph::element::i64, ngraph::element::i32 }});
 
+
     manager.run_passes(nGraphFunc);
 }
 

src/plugins/intel_cpu/tests/unit/ngraph_transformations/logsoftmax_decompose.cpp

@@ -0,0 +1,80 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <gtest/gtest.h>
+
+#include <string>
+#include <memory>
+
+#include <openvino/opsets/opset1.hpp>
+#include <openvino/opsets/opset8.hpp>
+#include <ngraph_transformations/op/fully_connected.hpp>
+#include <ngraph_transformations/convert_logsoftmax.hpp>
+#include <ngraph_transformations/fc_bias_fusion.hpp>
+#include <transformations/init_node_info.hpp>
+#include <transformations/utils/utils.hpp>
+#include <openvino/pass/manager.hpp>
+
+#include "common_test_utils/ngraph_test_utils.hpp"
+
+using namespace testing;
+using namespace ov::intel_cpu;
+
+TEST_F(TransformationTestsF, ConvertLogSoftmax1) {
+    int compute_axis = 1;
+    {
+        auto input = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::Shape{ 3, 2, 2 });
+        auto logsoftmax = std::make_shared<ov::opset8::LogSoftmax>(input, compute_axis);
+        function = std::make_shared<ov::Model>(ov::NodeVector{ logsoftmax }, ov::ParameterVector{ input });
+        manager.register_pass<ConvertLogSoftmax>();
+    }
+
+    {
+        auto input = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::Shape{ 3, 2, 2 });
+        auto axis = std::make_shared<ov::opset8::Constant>(ov::element::i64, ov::Shape{}, compute_axis);
+        auto xMax = std::make_shared<ov::opset8::ReduceMax>(input, axis, true);
+        auto subtract = std::make_shared<ov::opset8::Subtract>(input, xMax);
+        auto exp = std::make_shared<ov::opset8::Exp>(subtract);
+        auto s = std::make_shared<ov::opset8::ReduceSum>(exp, axis, true);
+        auto log = std::make_shared<ov::opset8::Log>(s);
+        auto result = std::make_shared<ov::opset8::Subtract>(subtract, log);
+        function_ref = std::make_shared<ov::Model>(ov::NodeVector{ result }, ov::ParameterVector{ input });
+    }
+}
+
+TEST_F(TransformationTestsF, ConvertLogSoftmaxSplit) {
+    int compute_axis = 1;
+    auto make_split = [](const std::shared_ptr<ov::Node>& input) {
+        auto split_axis = std::make_shared<ov::opset8::Constant>(ov::element::i64, ov::Shape{}, 1);
+        auto split_length = std::make_shared<ov::opset8::Constant>(ov::element::i64, ov::Shape{2}, std::vector<int64_t>{4, 6});
+        auto split = std::make_shared<ov::opset8::VariadicSplit>(input, split_axis, split_length);
+        return split;
+    };
+
+    {
+        auto input = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::Shape{ 2, 10 });
+        auto split = make_split(input);
+        auto logsoftmax = std::make_shared<ov::opset8::LogSoftmax>(split->output(1), compute_axis);
+        auto scale = std::make_shared<ov::opset8::Constant>(ov::element::f32, ov::Shape{}, 1);
+        auto multiply = std::make_shared<ov::opset8::Multiply>(split->output(0), scale);
+
+        function = std::make_shared<ov::Model>(ov::NodeVector{multiply, logsoftmax}, ov::ParameterVector{ input });
+        manager.register_pass<ConvertLogSoftmax>();
+    }
+
+    {
+        auto input = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::Shape{ 2, 10 });
+        auto split = make_split(input);
+        auto axis = std::make_shared<ov::opset8::Constant>(ov::element::i64, ov::Shape{}, compute_axis);
+        auto xMax = std::make_shared<ov::opset8::ReduceMax>(split->output(1), axis, true);
+        auto subtract = std::make_shared<ov::opset8::Subtract>(split->output(1), xMax);
+        auto exp = std::make_shared<ov::opset8::Exp>(subtract);
+        auto s = std::make_shared<ov::opset8::ReduceSum>(exp, axis, true);
+        auto log = std::make_shared<ov::opset8::Log>(s);
+        auto result = std::make_shared<ov::opset8::Subtract>(subtract, log);
+        auto scale = std::make_shared<ov::opset8::Constant>(ov::element::f32, ov::Shape{}, 1);
+        auto multiply = std::make_shared<ov::opset8::Multiply>(split->output(0), scale);
+        function_ref = std::make_shared<ov::Model>(ov::NodeVector{ multiply, result }, ov::ParameterVector{ input });
+    }
+}