XNNPACK: Route tanh GELU to approxgelu (pytorch#22851)

### Summary

The XNNPACK backend currently ignores GELU's `approximate` argument and
serializes both modes as `XNNGelu`. As a result,
`nn.GELU(approximate="tanh")` executes exact GELU after delegation. For
FP32 input `[-2.7]`, the baseline differs from the tanh reference by
approximately `4.73e-4`, failing comparison at `atol=rtol=1e-5`.

This PR serializes tanh GELU as `XNNApproxGelu` and dispatches it to
XNNPACK's existing `xnn_unary_approxgelu` operation. Default/exact GELU
and FP16 fallback remain unchanged. The new node is appended to both
FlatBuffer unions to preserve existing node IDs.

Affected XNNPACK models need to be re-exported, and the new node
requires an updated runtime.

### Test plan

I tested the fix on macOS arm64 using rebuilt runtimes based on
`500849ba5b`. Baseline validation reproduced the FP32 tanh GELU mismatch
at input `[-2.7]`, failing comparison at `atol=rtol=1e-5`. The current
parameterized regression tests retain this input and tolerance.

All 7 GELU tests passed, covering default GELU, both approximation
modes, static/dynamic shapes, and FP16 fallback. Local lintrunner
passed.

Authored with assistance from OpenAI Codex.


cc @GregoryComer @digantdesai @cbilgin @JakeStevens

Author: sqliangx

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/BUCK

@@ -94,7 +94,9 @@ fbcode_target(_kind = runtime.python_test,
     ]) + [
         "test_xnnpack_utils.py",
     ],
+    supports_static_listing = False,
     deps = [
+        "fbsource//third-party/pypi/parameterized:parameterized",
         "//executorch/backends/test/harness:tester",
         "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
         "//executorch/backends/xnnpack/quantizer:xnnpack_quantizer",

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,62 +31,90 @@ 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)
 
-    def run_gelu_test(self, inputs):
+    def run_gelu_test(
+        self, inputs, *, approximate="none", dynamic=False, atol=None, rtol=None
+    ):
         input_tensor = inputs[0]
+        is_fp16 = input_tensor.dtype == torch.float16
 
-        if input_tensor.dtype == torch.float16:
+        if is_fp16:
             with torch.no_grad():
                 ref_output = torch.nn.functional.gelu(
-                    input_tensor.to(torch.float32)
+                    input_tensor.to(torch.float32), approximate=approximate
                 ).to(torch.float16)
-            atol, rtol = calculate_fp16_gelu_tolerance(ref_output)
+            default_atol, default_rtol = calculate_fp16_gelu_tolerance(ref_output)
         else:
-            atol = 1e-03
-            rtol = 1e-03
+            default_atol, default_rtol = 1e-3, 1e-3
 
-        (
-            Tester(self.Gelu(), inputs)
+        if atol is None:
+            atol = default_atol
+        if rtol is None:
+            rtol = default_rtol
+
+        dynamic_shapes = (
+            ({0: torch.export.Dim("length", min=2, max=32)},) if dynamic else None
+        )
+        tester = (
+            Tester(
+                self.Gelu(approximate=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(atol=atol, rtol=rtol)
         )
 
-    def test_fp16_gelu(self):
-        # 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
-            )
-        atol, rtol = calculate_fp16_gelu_tolerance(ref_output)
+        if is_fp16:
+            # Older versions of XNNPACK don't support fp16 GELU.
+            # TODO (gjcomer) Remove this when we update XNNPACK. (#16679)
+            tester.check(
+                ["executorch_exir_dialects_edge__ops_aten_gelu_default"]
+            ).check_not(["torch.ops.higher_order.executorch_call_delegate"])
+        else:
+            tester.check_count(
+                {"torch.ops.higher_order.executorch_call_delegate": 1}
+            ).check_not(["executorch_exir_dialects_edge__ops_aten_gelu_default"])
 
         (
-            Tester(self.Gelu(), inputs)
-            .export()
-            .check_count({"torch.ops.aten.gelu.default": 1})
-            .to_edge_transform_and_lower()
-            # Expect no delegation
-            .check(["executorch_exir_dialects_edge__ops_aten_gelu_default"])
-            .check_not(["torch.ops.higher_order.executorch_call_delegate"])
-            .to_executorch()
+            tester.to_executorch()
             .serialize()
-            .run_method_and_compare_outputs(atol=atol, rtol=rtol)
+            .run_method_and_compare_outputs(inputs=inputs, atol=atol, rtol=rtol)
         )
+        if dynamic:
+            tester.run_method_and_compare_outputs(
+                inputs=(torch.linspace(-6, 6, 19),), atol=atol, rtol=rtol
+            )
+
+    @parameterized.expand([("none",), ("tanh",)])
+    def test_fp16_gelu(self, approximate):
+        inputs = (torch.randn(20).to(torch.float16),)
+        self.run_gelu_test(inputs, approximate=approximate)
 
     def test_fp32_gelu(self):
         inputs = (torch.randn(20),)
         self.run_gelu_test(inputs)
+
+    @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]),)
+        self.run_gelu_test(
+            inputs,
+            approximate=approximate,
+            dynamic=dynamic,
+            atol=1e-5,
+            rtol=1e-5,
+        )