fast tanh replace

src/include/migraphx/rewrite_gelu.hpp

@@ -34,7 +34,7 @@ inline namespace MIGRAPHX_INLINE_NS {
 struct module;
 
 /**
- * Rewrite gelu standard formula as the sigmoid approximation formula
+ * Rewrite GELU blocks as different approximations.
  */
 struct MIGRAPHX_EXPORT rewrite_gelu
 {

src/rewrite_gelu.cpp

@@ -26,11 +26,16 @@
 #include <migraphx/make_op.hpp>
 #include <migraphx/matcher.hpp>
 #include <migraphx/match/gelu_erf.hpp>
+#include <migraphx/match/gelu_tanh.hpp>
 #include <migraphx/common.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
+/**
+ * Finds GELU blocks using the Gaussian distribution and replaces them with the sigmoid approximation
+ * if the data type is fp16.
+ */
 struct find_gelu_erf
 {
     auto matcher() const { return match::gelu_erf(); }
@@ -53,7 +58,44 @@ struct find_gelu_erf
     }
 };
 
-void rewrite_gelu::apply(module& m) const { match::find_matches(m, find_gelu_erf{}); }
+/**
+ * Find fastGELU blocks (where the graph already does a GELU approximation) and replace them
+ * with an alternative approximation that is less likely to overflow.
+ * The replacement approximation is equivalent to:
+ * GELU(x) ~= 0.5 * x * ( 1 + tanh( sqrt(2/M_PI) * (x + 0.044715 * x^3)))
+ */
+struct find_tanh_fast_gelu
+{
+    auto matcher() const { return match::gelu_tanh(); }
+
+    void apply(module& m, const match::matcher_result& r) const
+    {
+        auto ins = r.result;
+        auto x   = r.instructions["x"];
+        auto sqrt_2_rpi = m.add_literal(literal{shape{x->get_shape().type()}, {0.7978845608028653558798921198687637369517172623298693153318516593}});
+        auto fit_const = m.add_literal(literal{shape{x->get_shape().type()}, {0.044715f}});
+        auto one = m.add_literal(literal{shape{x->get_shape().type()}, {1.0f}});
+        auto xb = insert_common_op(m, ins, make_op("mul"), {x, sqrt_2_rpi});
+        auto a = insert_common_op(m, ins, make_op("mul"), {xb, fit_const});
+        auto b = m.insert_instruction(ins, make_op("mul"), a, x);
+        auto c = m.insert_instruction(ins, make_op("mul"), b, x);
+        auto u = m.insert_instruction(ins, make_op("add"), c, xb);
+        auto neg_u = m.insert_instruction(ins, make_op("neg"), u);
+        auto d = m.insert_instruction(ins, make_op("sub"), neg_u, u);
+        auto emu = m.insert_instruction(ins, make_op("exp"), d);
+        auto e = insert_common_op(m, ins, make_op("add"), {one, emu});
+        auto cdf = insert_common_op(m, ins, make_op("div"), {one, e});
+        auto y = m.insert_instruction(ins, make_op("mul"), x, cdf);
+        m.replace_instruction(ins, y);
+    }
+};
+
+void rewrite_gelu::apply(module& m) const {
+    match::find_matches(m,
+            find_gelu_erf{},
+            find_tanh_fast_gelu{}  
+        );
+}
 
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx