diff --git a/backends/xnnpack/operators/op_gelu.py b/backends/xnnpack/operators/op_gelu.py index d96ddb0024e..aee72a744e6 100644 --- a/backends/xnnpack/operators/op_gelu.py +++ b/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, diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 2b8ab9b3181..1afea0ea9ef 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/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) diff --git a/backends/xnnpack/serialization/runtime_schema.fbs b/backends/xnnpack/serialization/runtime_schema.fbs index fb2c9b1598c..888384e032e 100644 --- a/backends/xnnpack/serialization/runtime_schema.fbs +++ b/backends/xnnpack/serialization/runtime_schema.fbs @@ -159,6 +159,7 @@ union XNodeUnion { XNNSin: _XNNNode1x1, XNNCopy: _XNNNode1x1, XNNCos: _XNNNode1x1, + XNNApproxGelu: _XNNNode1x1, } union XValueUnion { diff --git a/backends/xnnpack/serialization/schema.fbs b/backends/xnnpack/serialization/schema.fbs index 203469421d1..96e4256f078 100644 --- a/backends/xnnpack/serialization/schema.fbs +++ b/backends/xnnpack/serialization/schema.fbs @@ -155,6 +155,7 @@ union XNodeUnion { XNNSin: _XNNNode1x1, XNNCopy: _XNNNode1x1, XNNCos: _XNNNode1x1, + XNNApproxGelu: _XNNNode1x1, } union XValueUnion { diff --git a/backends/xnnpack/serialization/xnnpack_graph_schema.py b/backends/xnnpack/serialization/xnnpack_graph_schema.py index e95a55e1c01..360034c1987 100644 --- a/backends/xnnpack/serialization/xnnpack_graph_schema.py +++ b/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, ] diff --git a/backends/xnnpack/test/ops/test_gelu.py b/backends/xnnpack/test/ops/test_gelu.py index fdb7e2e2848..377b1d098d7 100644 --- a/backends/xnnpack/test/ops/test_gelu.py +++ b/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 + )