Add GELU fusion (#5067)

* GeLU fusion

* Fix matcher and sqrt(2) value

* Replace with math constant

* Correct comments in header file

* Add tests for transformations

* Add define of math constants in tests

* refactor code and add checks for nullptr

* Extend tests with failing examples

* change tests

inference-engine/src/transformations/include/transformations/common_optimizations/gelu_fusion.hpp

@@ -0,0 +1,68 @@
+// Copyright (C) 2018-2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <memory>
+#include <ngraph/pass/graph_rewrite.hpp>
+#include <transformations_visibility.hpp>
+#include <utility>
+
+namespace ngraph {
+namespace pass {
+
+class TRANSFORMATIONS_API GeluFusion;
+class TRANSFORMATIONS_API GeluFusionWithErfOne;
+class TRANSFORMATIONS_API GeluFusionWithErfTwo;
+class TRANSFORMATIONS_API GeluFusionWithErfThree;
+
+}  // namespace pass
+}  // namespace ngraph
+
+/**
+ * @ingroup ie_transformation_common_api
+ * @brief GeluFusion transformation replaces a sub-graph
+ * (0.5 * x) * (1 + erf(x / sqrt(2))) with a Gelu op.
+ */
+class ngraph::pass::GeluFusionWithErfOne : public ngraph::pass::MatcherPass {
+public:
+    NGRAPH_RTTI_DECLARATION;
+    GeluFusionWithErfOne();
+};
+
+/**
+ * @ingroup ie_transformation_common_api
+ * @brief GeluFusion transformation replaces a sub-graph
+ * 0.5 * (x * (1 + erf(x / sqrt(2))) with a Gelu op.
+ */
+class ngraph::pass::GeluFusionWithErfTwo : public ngraph::pass::MatcherPass {
+public:
+    NGRAPH_RTTI_DECLARATION;
+    GeluFusionWithErfTwo();
+};
+
+/**
+ * @ingroup ie_transformation_common_api
+ * @brief GeluFusion transformation replaces a sub-graph
+ * x * (0.5 * (1 + erf(x / sqrt(2))) with a Gelu op.
+ */
+class ngraph::pass::GeluFusionWithErfThree : public ngraph::pass::MatcherPass {
+public:
+    NGRAPH_RTTI_DECLARATION;
+    GeluFusionWithErfThree();
+};
+
+/**
+ * @ingroup ie_transformation_common_api
+ * @brief GeluFusion transformation replaces various sub-graphs with a Gelu op.
+ */
+class ngraph::pass::GeluFusion : public ngraph::pass::GraphRewrite {
+public:
+    NGRAPH_RTTI_DECLARATION;
+    GeluFusion() {
+        add_matcher<ngraph::pass::GeluFusionWithErfOne>();
+        add_matcher<ngraph::pass::GeluFusionWithErfTwo>();
+        add_matcher<ngraph::pass::GeluFusionWithErfThree>();
+    }
+};

inference-engine/src/transformations/src/transformations/common_optimizations/common_optimizations.cpp

@@ -13,6 +13,7 @@
 #include "transformations/common_optimizations/conv_mul_fusion.hpp"
 #include "transformations/common_optimizations/fq_mul_fusion.hpp"
 #include "transformations/common_optimizations/fq_reshape_fusion.hpp"
+#include "transformations/common_optimizations/gelu_fusion.hpp"
 #include "transformations/common_optimizations/depth_to_space_fusion.hpp"
 #include "transformations/common_optimizations/optimize_strided_slice.hpp"
 #include "transformations/common_optimizations/softplus_fusion.hpp"
@@ -113,6 +114,7 @@ bool ngraph::pass::CommonOptimizations::run_on_function(std::shared_ptr<ngraph::
     common_fusions->add_matcher<ngraph::pass::SpaceToBatchFusion>();
     common_fusions->add_matcher<ngraph::pass::BatchToSpaceFusion>();
     common_fusions->add_matcher<ngraph::pass::DilatedConvolutionConverter>();
+    common_fusions->add_matcher<ngraph::pass::GeluFusion>();
     common_fusions->set_name("ngraph::pass::CommonFusions");
 
     manager.register_pass<ngraph::pass::ConvertPadToGroupConvolution, false>();

inference-engine/src/transformations/src/transformations/common_optimizations/gelu_fusion.cpp

@@ -0,0 +1,229 @@
+// Copyright (C) 2018-2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#define _USE_MATH_DEFINES
+
+#include "transformations/common_optimizations/gelu_fusion.hpp"
+
+#include <math.h>
+
+#include <memory>
+#include <ngraph/opsets/opset7.hpp>
+#include <ngraph/pattern/op/wrap_type.hpp>
+#include <ngraph/rt_info.hpp>
+
+#include "itt.hpp"
+#include "transformations/utils/utils.hpp"
+
+NGRAPH_RTTI_DEFINITION(ngraph::pass::GeluFusion, "GeluFusion", 0);
+
+NGRAPH_RTTI_DEFINITION(ngraph::pass::GeluFusionWithErfOne,
+                       "GeluFusionWithErfOne", 0);
+
+ngraph::pass::GeluFusionWithErfOne::GeluFusionWithErfOne() {
+    MATCHER_SCOPE(GeluFusionWithErfOne);
+    // Replaces a sub-graph with a Gelu op
+    // Shared by every pattern: (1 + erf(x / sqrt(2)))
+    auto input = ngraph::pattern::any_input();
+    auto div_constant = ngraph::pattern::wrap_type<ngraph::opset7::Constant>();
+    auto div = ngraph::pattern::wrap_type<ngraph::opset7::Divide>(
+        {input, div_constant});
+    auto erf = ngraph::pattern::wrap_type<ngraph::opset7::Erf>({div});
+    auto add_constant = ngraph::pattern::wrap_type<ngraph::opset7::Constant>();
+    auto add =
+        ngraph::pattern::wrap_type<ngraph::opset7::Add>({add_constant, erf});
+    auto mul_constant = ngraph::pattern::wrap_type<ngraph::opset7::Constant>();
+
+    // (0.5 * x) * (1 + erf(x / sqrt(2))
+    auto mul_first = ngraph::pattern::wrap_type<ngraph::opset7::Multiply>(
+        {input, mul_constant});
+    auto mul =
+        ngraph::pattern::wrap_type<ngraph::opset7::Multiply>({mul_first, add});
+
+    ngraph::matcher_pass_callback callback = [=](ngraph::pattern::Matcher &m) {
+        auto &pattern_to_output = m.get_pattern_value_map();
+        auto x_output = pattern_to_output.at(input);
+
+        auto div_const_value =
+            std::dynamic_pointer_cast<ngraph::opset7::Constant>(
+                pattern_to_output.at(div_constant).get_node_shared_ptr());
+        auto add_const_value =
+            std::dynamic_pointer_cast<ngraph::opset7::Constant>(
+                pattern_to_output.at(add_constant).get_node_shared_ptr());
+        auto mul_const_value =
+            std::dynamic_pointer_cast<ngraph::opset7::Constant>(
+                pattern_to_output.at(mul_constant).get_node_shared_ptr());
+
+        if (!div_const_value || !add_const_value || !mul_const_value) {
+            return false;
+        }
+
+        bool valid_constant_values =
+            op::util::has_constant_value<float>(div_const_value, M_SQRT2) &&
+            op::util::has_constant_value<float>(add_const_value, 1.0f) &&
+            op::util::has_constant_value<float>(mul_const_value, 0.5f);
+
+        if (!valid_constant_values) {
+            return false;
+        }
+
+        auto gelu = std::make_shared<ngraph::opset7::Gelu>(x_output);
+
+        gelu->set_friendly_name(m.get_match_root()->get_friendly_name());
+        ngraph::copy_runtime_info(
+            {
+                pattern_to_output.at(div).get_node_shared_ptr(),
+                pattern_to_output.at(erf).get_node_shared_ptr(),
+                pattern_to_output.at(add).get_node_shared_ptr(),
+                pattern_to_output.at(mul_first).get_node_shared_ptr(),
+                pattern_to_output.at(mul).get_node_shared_ptr(),
+            },
+            gelu);
+        ngraph::replace_node(m.get_match_root(), gelu);
+        return true;
+    };
+
+    auto m = std::make_shared<ngraph::pattern::Matcher>(mul, matcher_name);
+    register_matcher(m, callback);
+}
+
+NGRAPH_RTTI_DEFINITION(ngraph::pass::GeluFusionWithErfTwo,
+                       "GeluFusionWithErfTwo", 0);
+
+ngraph::pass::GeluFusionWithErfTwo::GeluFusionWithErfTwo() {
+    MATCHER_SCOPE(GeluFusionWithErfTwo);
+    // Replaces a sub-graph with a Gelu op
+    // Shared by every pattern: (1 + erf(x / sqrt(2)))
+    auto input = ngraph::pattern::any_input();
+    auto div_constant = ngraph::pattern::wrap_type<ngraph::opset7::Constant>();
+    auto div = ngraph::pattern::wrap_type<ngraph::opset7::Divide>(
+        {input, div_constant});
+    auto erf = ngraph::pattern::wrap_type<ngraph::opset7::Erf>({div});
+    auto add_constant = ngraph::pattern::wrap_type<ngraph::opset7::Constant>();
+    auto add =
+        ngraph::pattern::wrap_type<ngraph::opset7::Add>({add_constant, erf});
+    auto mul_constant = ngraph::pattern::wrap_type<ngraph::opset7::Constant>();
+
+    // 0.5 * (x * (1 + erf(x / sqrt(2)))
+    auto mul_first =
+        ngraph::pattern::wrap_type<ngraph::opset7::Multiply>({input, add});
+    auto mul = ngraph::pattern::wrap_type<ngraph::opset7::Multiply>(
+        {mul_constant, mul_first});
+
+    ngraph::matcher_pass_callback callback = [=](ngraph::pattern::Matcher &m) {
+        auto &pattern_to_output = m.get_pattern_value_map();
+        auto x_output = pattern_to_output.at(input);
+
+        auto div_const_value =
+            std::dynamic_pointer_cast<ngraph::opset7::Constant>(
+                pattern_to_output.at(div_constant).get_node_shared_ptr());
+        auto add_const_value =
+            std::dynamic_pointer_cast<ngraph::opset7::Constant>(
+                pattern_to_output.at(add_constant).get_node_shared_ptr());
+        auto mul_const_value =
+            std::dynamic_pointer_cast<ngraph::opset7::Constant>(
+                pattern_to_output.at(mul_constant).get_node_shared_ptr());
+
+        if (!div_const_value || !add_const_value || !mul_const_value) {
+            return false;
+        }
+
+        bool valid_constant_values =
+            op::util::has_constant_value<float>(div_const_value, M_SQRT2) &&
+            op::util::has_constant_value<float>(add_const_value, 1.0f) &&
+            op::util::has_constant_value<float>(mul_const_value, 0.5f);
+
+        if (!valid_constant_values) {
+            return false;
+        }
+
+        auto gelu = std::make_shared<ngraph::opset7::Gelu>(x_output);
+
+        gelu->set_friendly_name(m.get_match_root()->get_friendly_name());
+        ngraph::copy_runtime_info(
+            {
+                pattern_to_output.at(div).get_node_shared_ptr(),
+                pattern_to_output.at(erf).get_node_shared_ptr(),
+                pattern_to_output.at(add).get_node_shared_ptr(),
+                pattern_to_output.at(mul_first).get_node_shared_ptr(),
+                pattern_to_output.at(mul).get_node_shared_ptr(),
+            },
+            gelu);
+        ngraph::replace_node(m.get_match_root(), gelu);
+        return true;
+    };
+
+    auto m = std::make_shared<ngraph::pattern::Matcher>(mul, matcher_name);
+    register_matcher(m, callback);
+}
+
+NGRAPH_RTTI_DEFINITION(ngraph::pass::GeluFusionWithErfThree,
+                       "GeluFusionWithErfThree", 0);
+
+ngraph::pass::GeluFusionWithErfThree::GeluFusionWithErfThree() {
+    MATCHER_SCOPE(GeluFusionWithErfThree);
+    // Replaces a sub-graph with a Gelu op
+    // Shared by every pattern: (1 + erf(x / sqrt(2)))
+    auto input = ngraph::pattern::any_input();
+    auto div_constant = ngraph::pattern::wrap_type<ngraph::opset7::Constant>();
+    auto div = ngraph::pattern::wrap_type<ngraph::opset7::Divide>(
+        {input, div_constant});
+    auto erf = ngraph::pattern::wrap_type<ngraph::opset7::Erf>({div});
+    auto add_constant = ngraph::pattern::wrap_type<ngraph::opset7::Constant>();
+    auto add =
+        ngraph::pattern::wrap_type<ngraph::opset7::Add>({add_constant, erf});
+    auto mul_constant = ngraph::pattern::wrap_type<ngraph::opset7::Constant>();
+
+    // x * (0.5 * (1 + erf(x / sqrt(2)))
+    auto mul_first = ngraph::pattern::wrap_type<ngraph::opset7::Multiply>(
+        {add, mul_constant});
+    auto mul = ngraph::pattern::wrap_type<ngraph::opset7::Multiply>(
+        {input, mul_first});
+
+    ngraph::matcher_pass_callback callback = [=](ngraph::pattern::Matcher &m) {
+        auto &pattern_to_output = m.get_pattern_value_map();
+        auto x_output = pattern_to_output.at(input);
+
+        auto div_const_value =
+            std::dynamic_pointer_cast<ngraph::opset7::Constant>(
+                pattern_to_output.at(div_constant).get_node_shared_ptr());
+        auto add_const_value =
+            std::dynamic_pointer_cast<ngraph::opset7::Constant>(
+                pattern_to_output.at(add_constant).get_node_shared_ptr());
+        auto mul_const_value =
+            std::dynamic_pointer_cast<ngraph::opset7::Constant>(
+                pattern_to_output.at(mul_constant).get_node_shared_ptr());
+
+        if (!div_const_value || !add_const_value || !mul_const_value) {
+            return false;
+        }
+
+        bool valid_constant_values =
+            op::util::has_constant_value<float>(div_const_value, M_SQRT2) &&
+            op::util::has_constant_value<float>(add_const_value, 1.0f) &&
+            op::util::has_constant_value<float>(mul_const_value, 0.5f);
+
+        if (!valid_constant_values) {
+            return false;
+        }
+
+        auto gelu = std::make_shared<ngraph::opset7::Gelu>(x_output);
+
+        gelu->set_friendly_name(m.get_match_root()->get_friendly_name());
+        ngraph::copy_runtime_info(
+            {
+                pattern_to_output.at(div).get_node_shared_ptr(),
+                pattern_to_output.at(erf).get_node_shared_ptr(),
+                pattern_to_output.at(add).get_node_shared_ptr(),
+                pattern_to_output.at(mul_first).get_node_shared_ptr(),
+                pattern_to_output.at(mul).get_node_shared_ptr(),
+            },
+            gelu);
+        ngraph::replace_node(m.get_match_root(), gelu);
+        return true;
+    };
+
+    auto m = std::make_shared<ngraph::pattern::Matcher>(mul, matcher_name);
+    register_matcher(m, callback);
+}

inference-engine/tests/functional/inference_engine/transformations/gelu_fusion.cpp

@@ -0,0 +1,212 @@
+// Copyright (C) 2018-2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#define _USE_MATH_DEFINES
+
+#include <gtest/gtest.h>
+#include <math.h>
+
+#include <memory>
+#include <ngraph/function.hpp>
+#include <ngraph/opsets/opset7.hpp>
+#include <ngraph/pass/constant_folding.hpp>
+#include <ngraph/pass/manager.hpp>
+#include <queue>
+#include <string>
+#include <transformations/common_optimizations/gelu_fusion.hpp>
+#include <transformations/init_node_info.hpp>
+#include <transformations/utils/utils.hpp>
+
+#include "common_test_utils/ngraph_test_utils.hpp"
+
+using namespace testing;
+using namespace ngraph;
+
+TEST(TransformationTests, GeluFusionPatternOne) {
+    std::shared_ptr<Function> f(nullptr), f_ref(nullptr);
+    {
+        auto data =
+            std::make_shared<opset7::Parameter>(element::f32, Shape{2, 2});
+
+        auto div_const =
+            opset7::Constant::create(element::f32, Shape{1}, {M_SQRT2});
+        auto add_const =
+            opset7::Constant::create(element::f32, Shape{1}, {1.0});
+        auto mul_const =
+            opset7::Constant::create(element::f32, Shape{1}, {0.5});
+
+        auto div = std::make_shared<opset7::Divide>(data, div_const);
+        auto erf = std::make_shared<opset7::Erf>(div);
+        auto add = std::make_shared<opset7::Add>(erf, add_const);
+        auto mul_first = std::make_shared<opset7::Multiply>(data, mul_const);
+        auto mul = std::make_shared<opset7::Multiply>(mul_first, add);
+
+        f = std::make_shared<Function>(NodeVector{mul}, ParameterVector{data});
+
+        pass::Manager m;
+        m.register_pass<pass::InitNodeInfo>();
+        m.register_pass<pass::GeluFusionWithErfOne>();
+        m.run_passes(f);
+        ASSERT_NO_THROW(check_rt_info(f));
+    }
+
+    {
+        auto data =
+            std::make_shared<opset1::Parameter>(element::f32, Shape{2, 2});
+        auto gelu = std::make_shared<opset7::Gelu>(data);
+        f_ref =
+            std::make_shared<Function>(NodeVector{gelu}, ParameterVector{data});
+    }
+
+    auto res = compare_functions(f, f_ref);
+    ASSERT_TRUE(res.first) << res.second;
+}
+
+TEST(TransformationTests, GeluFusionPatternTwo) {
+    std::shared_ptr<Function> f(nullptr), f_ref(nullptr);
+    {
+        auto data =
+            std::make_shared<opset7::Parameter>(element::f32, Shape{2, 2});
+
+        auto div_const =
+            opset7::Constant::create(element::f32, Shape{1}, {M_SQRT2});
+        auto add_const =
+            opset7::Constant::create(element::f32, Shape{1}, {1.0});
+        auto mul_const =
+            opset7::Constant::create(element::f32, Shape{1}, {0.5});
+
+        auto div = std::make_shared<opset7::Divide>(data, div_const);
+        auto erf = std::make_shared<opset7::Erf>(div);
+        auto add = std::make_shared<opset7::Add>(erf, add_const);
+        auto mul_first = std::make_shared<opset7::Multiply>(data, add);
+        auto mul = std::make_shared<opset7::Multiply>(mul_first, mul_const);
+
+        f = std::make_shared<Function>(NodeVector{mul}, ParameterVector{data});
+
+        pass::Manager m;
+        m.register_pass<pass::InitNodeInfo>();
+        m.register_pass<pass::GeluFusionWithErfTwo>();
+        m.run_passes(f);
+        ASSERT_NO_THROW(check_rt_info(f));
+    }
+
+    {
+        auto data =
+            std::make_shared<opset1::Parameter>(element::f32, Shape{2, 2});
+        auto gelu = std::make_shared<opset7::Gelu>(data);
+        f_ref =
+            std::make_shared<Function>(NodeVector{gelu}, ParameterVector{data});
+    }
+
+    auto res = compare_functions(f, f_ref);
+    ASSERT_TRUE(res.first) << res.second;
+}
+
+TEST(TransformationTests, GeluFusionPatternThree) {
+    std::shared_ptr<Function> f(nullptr), f_ref(nullptr);
+    {
+        auto data =
+            std::make_shared<opset7::Parameter>(element::f32, Shape{2, 2});
+
+        auto div_const =
+            opset7::Constant::create(element::f32, Shape{1}, {M_SQRT2});
+        auto add_const =
+            opset7::Constant::create(element::f32, Shape{1}, {1.0});
+        auto mul_const =
+            opset7::Constant::create(element::f32, Shape{1}, {0.5});
+
+        auto div = std::make_shared<opset7::Divide>(data, div_const);
+        auto erf = std::make_shared<opset7::Erf>(div);
+        auto add = std::make_shared<opset7::Add>(erf, add_const);
+        auto mul_first = std::make_shared<opset7::Multiply>(add, mul_const);
+        auto mul = std::make_shared<opset7::Multiply>(data, mul_first);
+
+        f = std::make_shared<Function>(NodeVector{mul}, ParameterVector{data});
+
+        pass::Manager m;
+        m.register_pass<pass::InitNodeInfo>();
+        m.register_pass<pass::GeluFusionWithErfThree>();
+        m.run_passes(f);
+        ASSERT_NO_THROW(check_rt_info(f));
+    }
+
+    {
+        auto data =
+            std::make_shared<opset1::Parameter>(element::f32, Shape{2, 2});
+        auto gelu = std::make_shared<opset7::Gelu>(data);
+        f_ref =
+            std::make_shared<Function>(NodeVector{gelu}, ParameterVector{data});
+    }
+
+    auto res = compare_functions(f, f_ref);
+    ASSERT_TRUE(res.first) << res.second;
+}
+
+TEST(TransformationTests, GeluFusionPatternIncorrectDivConstValue) {
+    std::shared_ptr<Function> f(nullptr), f_ref(nullptr);
+    {
+        auto data =
+            std::make_shared<opset7::Parameter>(element::f32, Shape{2, 2});
+
+        auto div_const =
+            opset7::Constant::create(element::f32, Shape{1}, {1.4149});
+        auto add_const =
+            opset7::Constant::create(element::f32, Shape{1}, {1.0});
+        auto mul_const =
+            opset7::Constant::create(element::f32, Shape{1}, {0.5});
+
+        auto div = std::make_shared<opset7::Divide>(data, div_const);
+        auto erf = std::make_shared<opset7::Erf>(div);
+        auto add = std::make_shared<opset7::Add>(erf, add_const);
+        auto mul_first = std::make_shared<opset7::Multiply>(data, add);
+        auto mul = std::make_shared<opset7::Multiply>(mul_first, mul_const);
+
+        f = std::make_shared<Function>(NodeVector{mul}, ParameterVector{data});
+        f_ref =
+            std::make_shared<Function>(NodeVector{mul}, ParameterVector{data});
+
+        pass::Manager m;
+        m.register_pass<pass::InitNodeInfo>();
+        m.register_pass<pass::GeluFusionWithErfTwo>();
+        m.run_passes(f);
+        ASSERT_NO_THROW(check_rt_info(f));
+    }
+
+    auto res = compare_functions(f, f_ref);
+    ASSERT_TRUE(res.first) << res.second;
+}
+
+TEST(TransformationTests, GeluFusionPatternTooShortDivConstValue) {
+    std::shared_ptr<Function> f(nullptr), f_ref(nullptr);
+    {
+        auto data =
+            std::make_shared<opset7::Parameter>(element::f32, Shape{2, 2});
+
+        auto div_const =
+            opset7::Constant::create(element::f32, Shape{1}, {1.4142});
+        auto add_const =
+            opset7::Constant::create(element::f32, Shape{1}, {1.0});
+        auto mul_const =
+            opset7::Constant::create(element::f32, Shape{1}, {0.5});
+
+        auto div = std::make_shared<opset7::Divide>(data, div_const);
+        auto erf = std::make_shared<opset7::Erf>(div);
+        auto add = std::make_shared<opset7::Add>(erf, add_const);
+        auto mul_first = std::make_shared<opset7::Multiply>(data, add);
+        auto mul = std::make_shared<opset7::Multiply>(mul_first, mul_const);
+
+        f = std::make_shared<Function>(NodeVector{mul}, ParameterVector{data});
+        f_ref =
+            std::make_shared<Function>(NodeVector{mul}, ParameterVector{data});
+
+        pass::Manager m;
+        m.register_pass<pass::InitNodeInfo>();
+        m.register_pass<pass::GeluFusionWithErfTwo>();
+        m.run_passes(f);
+        ASSERT_NO_THROW(check_rt_info(f));
+    }
+
+    auto res = compare_functions(f, f_ref);
+    ASSERT_TRUE(res.first) << res.second;
+}