Add Gelu Tanh fusion rule

onnxscript/rewriter/ort_fusions/gelu.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import math
+
+from onnxscript import ir
+from onnxscript.rewriter import pattern
+
+_sqrt_two_over_pi = math.sqrt(2.0 / math.pi)
+
+
+class GeluTanhFusion(pattern.RewriteRuleClassBase):
+    def pattern(self, op, x):
+        # GELU(x) = 0.5 * x * {1 + Tanh[\sqrt(2/pi) * (x + 0.044715 * x^3)]}
+        t1 = op.Pow(x, 3)
+        t2 = op.Mul(0.044715, t1)
+        t3 = op.Add(x, t2)
+
+        t4 = op.Mul(_sqrt_two_over_pi, t3)
+        t5 = op.Tanh(t4)
+        t6 = op.Add(t5, 1)
+        t7 = op.Mul(x, t6)
+        result = op.Mul(0.5, t7)
+        return result
+
+    def rewrite(self, op, x):
+        return op.Gelu(x, _domain="com.microsoft")
+
+
+_rule = GeluTanhFusion.rule()
+
+gelu_rules = pattern.RewriteRuleSet([_rule])
+
+
+def fuse_gelu(model: ir.Model) -> None:
+    gelu_rules.apply_to_model(model)

onnxscript/rewriter/ort_fusions/gelu_test.py

@@ -0,0 +1,57 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import math
+import unittest
+
+import numpy as np
+
+import onnxscript.ir as ir
+import onnxscript.rewriter.ort_fusions._test_utils as test_utils
+from onnxscript import FLOAT, script
+from onnxscript import opset18 as op
+from onnxscript.optimizer import optimize, remove_unused_nodes
+from onnxscript.rewriter.ort_fusions.gelu import fuse_gelu
+
+
+class GeluFusionTest(unittest.TestCase):
+    def test_gelu_fusion(self):
+        _sqrt_two_over_pi = math.sqrt(2.0 / math.pi)
+
+        @script()
+        def gelu_model(x):
+            # GELU(x) = 0.5 * x * {1 + Tanh[\sqrt(2/pi) * (x + 0.044715 * x^3)]}
+            t1 = op.Pow(x, 3)
+            t2 = op.Mul(0.044715, t1)
+            t3 = op.Add(x, t2)
+
+            t4 = op.Mul(_sqrt_two_over_pi, t3)
+            t5 = op.Tanh(t4)
+            t6 = op.Add(t5, 1)
+            t7 = op.Mul(x, t6)
+            result = op.Mul(0.5, t7)
+            return result
+
+        model_proto = gelu_model.to_model_proto(
+            input_types=[FLOAT[10]], output_types=[FLOAT[10]]
+        )
+        model = ir.serde.deserialize_model(model_proto)
+
+        # Eliminate redundant CastLike ops:
+        optimize(model)
+
+        input = {"x": np.random.randn(10).astype(np.float32)}
+        original_output = test_utils.ort_run("Original", model, input)
+
+        fuse_gelu(model)
+        remove_unused_nodes(model)
+
+        self.assertEqual(len(model.graph), 1)
+        self.assertEqual(model.graph.node(0).op_type, "Gelu")
+
+        optimized_output = test_utils.ort_run("Optimized", model, input)
+        test_utils.assert_allclose(original_output, optimized_output)
+
+
+if __name__ == "__main__":
+    unittest.main()