XNNPACK: Route tanh GELU to approxgelu

backends/xnnpack/operators/op_gelu.py

@@ -12,6 +12,7 @@
     register_node_visitor,
 )
 from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
+    XNNApproxGelu,
     XNNGelu,
     XNNGraph,
     XNode,
@@ -41,8 +42,15 @@ def define_node(
         # output
         output_id = vals_to_ids[node]
 
+        approximate = node.kwargs.get("approximate", "none")
+        if approximate == "none":
+            gelu_node_type = XNNGelu
+        elif approximate == "tanh":
+            gelu_node_type = XNNApproxGelu
+        else:
+            raise ValueError(f"Unsupported GELU approximation: {approximate}")
         ser_node = XNode(
-            xnode_union=XNNGelu(
+            xnode_union=gelu_node_type(
                 input_id=input_id,
                 output_id=output_id,
                 flags=0,

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -1980,6 +1980,7 @@ _DEFINE_UNARY_NODE_NO_PARAMS(
     xnn_unary_reciprocal_square_root)
 _DEFINE_UNARY_NODE_NO_PARAMS(Ceiling, xnn_unary_ceiling)
 _DEFINE_UNARY_NODE_NO_PARAMS(Gelu, xnn_unary_gelu)
+_DEFINE_UNARY_NODE_NO_PARAMS(ApproxGelu, xnn_unary_approxgelu)
 _DEFINE_UNARY_NODE_NO_PARAMS(Hardswish, xnn_unary_hardswish)
 _DEFINE_UNARY_NODE_NO_PARAMS(Log, xnn_unary_log)
 _DEFINE_UNARY_NODE_NO_PARAMS(Negate, xnn_unary_negate)
@@ -2021,6 +2022,7 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) {
     _DEFINE(ReciprocalSquareRoot)
     _DEFINE(Ceiling)
     _DEFINE(Gelu)
+    _DEFINE(ApproxGelu)
     _DEFINE(Hardswish)
     _DEFINE(Log)
     _DEFINE(Tanh)

backends/xnnpack/serialization/runtime_schema.fbs

@@ -159,6 +159,7 @@ union XNodeUnion {
   XNNSin: _XNNNode1x1,
   XNNCopy: _XNNNode1x1,
   XNNCos: _XNNNode1x1,
+  XNNApproxGelu: _XNNNode1x1,
 }
 
 union XValueUnion {

backends/xnnpack/serialization/schema.fbs

@@ -155,6 +155,7 @@ union XNodeUnion {
   XNNSin: _XNNNode1x1,
   XNNCopy: _XNNNode1x1,
   XNNCos: _XNNNode1x1,
+  XNNApproxGelu: _XNNNode1x1,
 }
 
 union XValueUnion {

backends/xnnpack/serialization/xnnpack_graph_schema.py

@@ -301,6 +301,11 @@ class XNNGelu(XNNNode1x1):
     pass
 
 
+@dataclass
+class XNNApproxGelu(XNNNode1x1):
+    pass
+
+
 @dataclass
 class XNNHardswish(XNNNode1x1):
     pass
@@ -421,6 +426,7 @@ class XNNScaledDotProductAttention:
     XNNSin,
     XNNCopy,
     XNNCos,
+    XNNApproxGelu,
 ]
 
 

backends/xnnpack/test/ops/test_gelu.py

@@ -8,6 +8,7 @@
 
 import torch
 from executorch.backends.xnnpack.test.tester import Tester
+from parameterized import parameterized
 
 
 def calculate_fp16_gelu_tolerance(ref_output_tensor):
@@ -30,9 +31,9 @@ def setUp(self):
         torch._dynamo.reset()
 
     class Gelu(torch.nn.Module):
-        def __init__(self):
+        def __init__(self, approximate="none"):
             super().__init__()
-            self.gelu = torch.nn.GELU()
+            self.gelu = torch.nn.GELU(approximate=approximate)
 
         def forward(self, x):
             return self.gelu(x)
@@ -62,19 +63,20 @@ def run_gelu_test(self, inputs):
             .run_method_and_compare_outputs(atol=atol, rtol=rtol)
         )
 
-    def test_fp16_gelu(self):
+    @parameterized.expand([("none",), ("tanh",)])
+    def test_fp16_gelu(self, approximate):
         # Older versions of XNNPACK don't support fp16 GELU.
         # TODO (gjcomer) Remove this when we update XNNPACK. (#16679)
         inputs = (torch.randn(20).to(torch.float16),)
 
         with torch.no_grad():
-            ref_output = torch.nn.functional.gelu(inputs[0].to(torch.float32)).to(
-                torch.float16
-            )
+            ref_output = torch.nn.functional.gelu(
+                inputs[0].to(torch.float32), approximate=approximate
+            ).to(torch.float16)
         atol, rtol = calculate_fp16_gelu_tolerance(ref_output)
 
         (
-            Tester(self.Gelu(), inputs)
+            Tester(self.Gelu(approximate=approximate), inputs)
             .export()
             .check_count({"torch.ops.aten.gelu.default": 1})
             .to_edge_transform_and_lower()
@@ -83,9 +85,51 @@ def test_fp16_gelu(self):
             .check_not(["torch.ops.higher_order.executorch_call_delegate"])
             .to_executorch()
             .serialize()
-            .run_method_and_compare_outputs(atol=atol, rtol=rtol)
+            .run_method_and_compare_outputs(inputs=inputs, atol=atol, rtol=rtol)
         )
 
     def test_fp32_gelu(self):
         inputs = (torch.randn(20),)
         self.run_gelu_test(inputs)
+
+    def test_fp32_gelu_tanh(self):
+        inputs = (torch.tensor([-2.7]),)
+        (
+            Tester(self.Gelu(approximate="tanh"), inputs)
+            .export()
+            .check_count({"torch.ops.aten.gelu.default": 1})
+            .to_edge_transform_and_lower()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .check_not(["executorch_exir_dialects_edge__ops_aten_gelu_default"])
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs(inputs=inputs, atol=1e-5, rtol=1e-5)
+        )
+
+    @parameterized.expand(
+        [
+            (approximate, dynamic)
+            for approximate in ("none", "tanh")
+            for dynamic in (False, True)
+        ]
+    )
+    def test_fp32_gelu_approximation(self, approximate, dynamic):
+        inputs = (torch.tensor([-6.0, -2.7, -1.0, 0.0, 1.0, 2.7, 6.0]),)
+        dynamic_shapes = (
+            ({0: torch.export.Dim("length", min=2, max=32)},) if dynamic else None
+        )
+        tester = (
+            Tester(self.Gelu(approximate), inputs, dynamic_shapes=dynamic_shapes)
+            .export()
+            .check_count({"torch.ops.aten.gelu.default": 1})
+            .to_edge_transform_and_lower()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .check_not(["executorch_exir_dialects_edge__ops_aten_gelu_default"])
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs(inputs=inputs, atol=1e-5, rtol=1e-5)
+        )
+        if dynamic:
+            tester.run_method_and_compare_outputs(
+                inputs=(torch.linspace(-6, 6, 19),), atol=1e-5, rtol=1e-5
+            )