Move gelu matcher to its own pass and add test

src/CMakeLists.txt

@@ -80,6 +80,7 @@ add_library(migraphx
     replace_allocate.cpp
     simplify_qdq.cpp
     rewrite_batchnorm.cpp
+    rewrite_gelu.cpp
     rewrite_pooling.cpp
     rewrite_quantization.cpp
     rewrite_rnn.cpp

src/include/migraphx/rewrite_gelu.hpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#ifndef MIGRAPHX_GUARD_RTGLIB_REWRITE_GELU_HPP
+#define MIGRAPHX_GUARD_RTGLIB_REWRITE_GELU_HPP
+
+#include <string>
+#include <migraphx/instruction_ref.hpp>
+#include <migraphx/config.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+struct module;
+
+/**
+ * Rewrite gelu standard formula as the sigmoid approximation formula
+ */
+struct rewrite_gelu
+{
+    std::string name() const { return "rewrite_gelu"; }
+    void apply(module& m) const;
+};
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/rewrite_gelu.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include <migraphx/rewrite_gelu.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/matcher.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+struct find_gelu_erf
+{
+    static auto match_div()
+    {
+        return match::name("div")(
+            match::either_arg(0, 1)(match::any().bind("x"), match::skip_broadcasts(match::has_value(1.414f, 1e-3))));
+    }
+
+    static auto match_erf() { return match::name("erf")(match::arg(0)(match_div())); }
+
+    static auto match_add()
+    {
+        return match::name("add")(
+            match::either_arg(0, 1)(match_erf(), match::skip_broadcasts(match::has_value(1.0f))));
+    }
+
+    static auto match_mul() { return match::name("mul")(match::either_arg(0, 1)(match::any(), match_add())); }
+
+    auto matcher() const
+    {
+        return match::name("mul")(
+            match::either_arg(0, 1)(match_mul(), match::skip_broadcasts(match::has_value(0.5f))));
+    }
+
+    void apply(module& m, const match::matcher_result& r) const
+    {
+        auto ins = r.result;
+        auto x   = r.instructions["x"];
+
+        auto lit = m.add_literal(literal{shape{x->get_shape().type()}, {1.702f}});
+        auto mul = m.insert_instruction(
+            ins, make_op("multibroadcast", {{"out_lens", x->get_shape().lens()}}), lit);
+        mul      = m.insert_instruction(ins, make_op("mul"), x, mul);
+        auto sig = m.insert_instruction(ins, make_op("neg"), mul);
+        sig = m.insert_instruction(ins, make_op("exp"), sig);
+        auto one = m.add_literal(literal{shape{x->get_shape().type()}, {1.0f}});
+        one      = m.insert_instruction(
+            ins, make_op("multibroadcast", {{"out_lens", x->get_shape().lens()}}), one);
+        sig = m.insert_instruction(ins, make_op("add"), sig, one);
+        sig = m.insert_instruction(ins, make_op("div"), one, sig);
+        sig = m.insert_instruction(ins, make_op("mul"), x, sig);
+        m.replace_instruction(ins, sig);
+    }
+};
+
+void rewrite_gelu::apply(module& m) const
+{
+    match::find_matches(m, find_gelu_erf{});
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/simplify_algebra.cpp

@@ -851,63 +851,6 @@ struct find_div_const
     }
 };
 
-struct find_gelu_erf
-{
-    static auto match_div()
-    {
-        return match::name("div")(
-            args(match::any().bind("x"), match::skip_broadcasts(match::is_constant())));
-    }
-
-    static auto match_mul1()
-    {
-        return match::name("mul")(
-            args(match::any().bind("x"), match::skip_broadcasts(match::name("recip"))));
-    }
-
-    static auto match_erf() { return match::name("erf")(match::arg(0)(match_div())); }
-
-    static auto match_add2()
-    {
-        return match::name("add")(
-            args(match_erf(), match::skip_broadcasts(match::has_value(1.0f))));
-    }
-
-    static auto match_add1()
-    {
-        return match::name("add")(
-            args(match::skip_broadcasts(match::is_constant()), match::name("dot")));
-    }
-
-    static auto match_mul2() { return match::name("mul")(match::args(match_add1(), match_add2())); }
-
-    auto matcher() const
-    {
-        return match::name("mul")(
-            match::args(match_mul2(), match::skip_broadcasts(match::has_value(0.5f))));
-    }
-
-    void apply(module& m, const match::matcher_result& r) const
-    {
-        auto ins = r.result;
-        auto x   = r.instructions["x"];
-
-        auto lit = m.add_literal(literal{shape{x->get_shape().type()}, {1.702f}});
-        auto mul = m.insert_instruction(
-            ins, make_op("multibroadcast", {{"out_lens", x->get_shape().lens()}}), lit);
-        mul      = m.insert_instruction(ins, make_op("mul"), x, mul);
-        auto sig = m.insert_instruction(ins, make_op("neg"), mul);
-        auto one = m.add_literal(literal{shape{x->get_shape().type()}, {1.0f}});
-        one      = m.insert_instruction(
-            ins, make_op("multibroadcast", {{"out_lens", x->get_shape().lens()}}), one);
-        sig = m.insert_instruction(ins, make_op("exp"), sig);
-        sig = m.insert_instruction(ins, make_op("add"), one, sig);
-        sig = m.insert_instruction(ins, make_op("div"), one, sig);
-        sig = m.insert_instruction(ins, make_op("mul"), x, sig);
-        m.replace_instruction(ins, sig);
-    }
-};
-
 struct find_sub_const
 {
     auto matcher() const
@@ -1106,7 +1049,6 @@ void simplify_algebra::apply(module& m) const
                             find_mul_slice_conv{},
                             find_mul_add{},
                             find_div_const{},
-                            find_gelu_erf{},
                             find_sub_const{},
                             find_rsqrt{},
                             find_concat_op{},

src/targets/gpu/target.cpp

@@ -42,6 +42,7 @@
 #include <migraphx/register_target.hpp>
 #include <migraphx/replace_allocate.hpp>
 #include <migraphx/rewrite_batchnorm.hpp>
+#include <migraphx/rewrite_gelu.hpp>
 #include <migraphx/rewrite_pooling.hpp>
 #include <migraphx/rewrite_quantization.hpp>
 #include <migraphx/rewrite_rnn.hpp>
@@ -115,6 +116,8 @@ std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_opti
         inline_module{},
         rewrite_pooling{},
         dead_code_elimination{},
+        rewrite_gelu{},
+        dead_code_elimination{},
         eliminate_common_subexpression{},
         dead_code_elimination{},
         simplify_algebra{},

test/rewrite_gelu_test.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <migraphx/rewrite_gelu.hpp>
+#include <migraphx/dead_code_elimination.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/ref/target.hpp>
+#include <migraphx/op/convolution.hpp>
+#include <migraphx/op/reshape.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/ranges.hpp>
+#include <test.hpp>
+#include <migraphx/make_op.hpp>
+
+#include <migraphx/serialize.hpp>
+
+#include <migraphx/verify.hpp>
+
+
+TEST_CASE(bias_gelu)
+{
+    migraphx::shape s1{migraphx::shape::half_type, {2, 4, 8}};
+    migraphx::shape s2{migraphx::shape::half_type};
+    migraphx::module m1;
+    {
+        auto a    = m1.add_parameter("a", s1);
+        auto b    = m1.add_parameter("b", s1);
+        auto add1 = m1.add_instruction(migraphx::make_op("add"), a, b);
+        auto l1   = m1.add_literal(migraphx::literal{s2, {1.4140625f}});
+        l1 = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l1);
+        auto div = m1.add_instruction(migraphx::make_op("div"), add1, l1);
+        auto erf = m1.add_instruction(migraphx::make_op("erf"), div);
+        auto l2  = m1.add_literal(migraphx::literal{s2, {1.0f}});
+        l2 = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l2);
+        auto add2 = m1.add_instruction(migraphx::make_op("add"), erf, l2);
+        auto mul  = m1.add_instruction(migraphx::make_op("mul"), add1, add2);
+        auto l3   = m1.add_literal(migraphx::literal{s2, {0.5f}});
+        l3 = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l3);
+        mul = m1.add_instruction(migraphx::make_op("mul"), mul, l3);
+        m1.add_return({mul});
+    }
+    migraphx::rewrite_gelu pass;
+    pass.apply(m1);
+    migraphx::dead_code_elimination dce;
+    dce.apply(m1);
+
+    migraphx::module m2;
+    {
+        auto a   = m2.add_parameter("a", s1);
+        auto b   = m2.add_parameter("b", s1);
+        auto add = m2.add_instruction(migraphx::make_op("add"), a, b);
+        auto l1  = m2.add_literal(migraphx::literal{s2, {1.702f}});
+        l1 = m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l1);
+        auto mul = m2.add_instruction(migraphx::make_op("mul"), add, l1);
+        auto sig = m2.add_instruction(migraphx::make_op("neg"), mul);
+        sig      = m2.add_instruction(migraphx::make_op("exp"), sig);
+        auto l2  = m2.add_literal(migraphx::literal{s2, {1.0f}});
+        l2 = m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l2);
+        sig = m2.add_instruction(migraphx::make_op("add"), sig, l2);
+        sig = m2.add_instruction(migraphx::make_op("div"), l2, sig);
+        sig = m2.add_instruction(migraphx::make_op("mul"), add, sig);
+        m2.add_return({sig});
+    }
+
+    EXPECT(m1 == m2);
+}
+
+TEST_CASE(non_bias_gelu)
+{
+    migraphx::shape s1{migraphx::shape::half_type, {2, 4, 8}};
+    migraphx::shape s2{migraphx::shape::half_type};
+    migraphx::module m1;
+    {
+        auto a    = m1.add_parameter("a", s1);
+        auto b    = m1.add_parameter("b", s1);
+        auto sub = m1.add_instruction(migraphx::make_op("sub"), a, b);
+        auto l1   = m1.add_literal(migraphx::literal{s2, {1.4140625f}});
+        l1 = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l1);
+        auto div = m1.add_instruction(migraphx::make_op("div"), sub, l1);
+        auto erf = m1.add_instruction(migraphx::make_op("erf"), div);
+        auto l2  = m1.add_literal(migraphx::literal{s2, {1.0f}});
+        l2 = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l2);
+        auto add2 = m1.add_instruction(migraphx::make_op("add"), erf, l2);
+        auto mul  = m1.add_instruction(migraphx::make_op("mul"), sub, add2);
+        auto l3   = m1.add_literal(migraphx::literal{s2, {0.5f}});
+        l3 = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l3);
+        mul = m1.add_instruction(migraphx::make_op("mul"), mul, l3);
+        m1.add_return({mul});
+    }
+    migraphx::rewrite_gelu pass;
+    pass.apply(m1);
+    migraphx::dead_code_elimination dce;
+    dce.apply(m1);
+
+    migraphx::module m2;
+    {
+        auto a   = m2.add_parameter("a", s1);
+        auto b   = m2.add_parameter("b", s1);
+        auto sub = m2.add_instruction(migraphx::make_op("sub"), a, b);
+        auto l1  = m2.add_literal(migraphx::literal{s2, {1.702f}});
+        l1 = m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l1);
+        auto mul = m2.add_instruction(migraphx::make_op("mul"), sub, l1);
+        auto sig = m2.add_instruction(migraphx::make_op("neg"), mul);
+        sig      = m2.add_instruction(migraphx::make_op("exp"), sig);
+        auto l2  = m2.add_literal(migraphx::literal{s2, {1.0f}});
+        l2 = m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l2);
+        sig = m2.add_instruction(migraphx::make_op("add"), sig, l2);
+        sig = m2.add_instruction(migraphx::make_op("div"), l2, sig);
+        sig = m2.add_instruction(migraphx::make_op("mul"), sub, sig);
+        m2.add_return({sig});
+    }
+
+    EXPECT(m1 == m2);
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/simplify_algebra_test.cpp

@@ -734,51 +734,6 @@ TEST_CASE(simplify_div_const)
     EXPECT(m1 == m2);
 }
 
-TEST_CASE(simplify_gelu_erf_with_bias)
-{
-    migraphx::shape s1{migraphx::shape::half_type, {2, 4, 8}};
-    migraphx::shape s2{migraphx::shape::half_type, {1}};
-    migraphx::module m1;
-    {
-        auto a    = m1.add_parameter("a", s1);
-        auto b    = m1.add_parameter("b", s1);
-        auto add1 = m1.add_instruction(migraphx::make_op("add"), a, b);
-        auto l1   = m1.add_literal(migraphx::literal{s2, {1.4140625f}});
-        l1 = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l1);
-        auto div = m1.add_instruction(migraphx::make_op("div"), add1, l1);
-        auto erf = m1.add_instruction(migraphx::make_op("erf"), div);
-        auto l2  = m1.add_literal(migraphx::literal{s2, {1.0f}});
-        l2 = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l2);
-        auto add2 = m1.add_instruction(migraphx::make_op("add"), erf, l2);
-        auto mul  = m1.add_instruction(migraphx::make_op("mul"), add1, add2);
-        auto l3   = m1.add_literal(migraphx::literal{s2, {0.5f}});
-        l3 = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l3);
-        m1.add_instruction(migraphx::make_op("mul"), mul, l3);
-    }
-    run_pass(m1);
-    run_pass(m1);
-    run_pass(m1);
-
-    migraphx::module m2;
-    {
-        auto a   = m2.add_parameter("a", s1);
-        auto b   = m2.add_parameter("b", s1);
-        auto add = m2.add_instruction(migraphx::make_op("add"), a, b);
-        auto l1  = m2.add_literal(migraphx::literal{s2, {1.702f}});
-        l1 = m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l1);
-        auto mul = m2.add_instruction(migraphx::make_op("mul"), add, l1);
-        auto sig = m2.add_instruction(migraphx::make_op("neg"), mul);
-        sig      = m2.add_instruction(migraphx::make_op("exp"), sig);
-        auto l2  = m2.add_literal(migraphx::literal{s2, {1.0f}});
-        l2 = m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s1.lens()}}), l2);
-        sig = m2.add_instruction(migraphx::make_op("add"), sig, l2);
-        sig = m2.add_instruction(migraphx::make_op("div"), l2, sig);
-        m2.add_instruction(migraphx::make_op("mul"), sig, add);
-    }
-
-    EXPECT(m1 == m2);
-}
-
 TEST_CASE(simplify_sub_const)
 {
     migraphx::module m1;