Support approximate gelu (#11246)

Summary:
Pull Request resolved: #11246

GELU accepts an `approximate` argument which is either `none` by default, or `tanh`

When the `approximate` kwarg is present, decompose the op.

We already have an existing test in test_aten_gelu_out to make sure the op is supported.

Reviewed By: zonglinpeng

Differential Revision: D75454999

Author: Michael Maitland

backends/cadence/aot/replace_ops.py

@@ -2065,11 +2065,10 @@ def call_operator(
         return super().call_operator(op, args, kwargs, meta)
 
 
-@register_cadence_pass(CadencePassAttribute(opt_level=2))
-class ReplaceGeluWithApproximateGeluPass(ExportPass):
+@register_cadence_pass(CadencePassAttribute(opt_level=0))
+class ReplaceAtenApproxGeluWithApproxGeluPass(ExportPass):
     """
-    Replace the gelu op with an approximate gelu op. The approximate gelu op
-    is more efficient on DSP backends.
+    Replace the aten gelu op with an approximate arg with an approximate gelu op.
     """
 
     def call_operator(
@@ -2079,6 +2078,9 @@ def call_operator(
         kwargs: Dict[str, Argument],
         meta: NodeMetadata,
     ) -> ProxyValue:
+        if "approximate" not in kwargs:
+            return super().call_operator(op, args, kwargs, meta)
+
         if op not in {
             exir_ops.edge.aten.gelu.default,
         }:
@@ -2414,7 +2416,7 @@ class CadenceReplaceOpsInGraph:
         ReplaceSingleElementTensorArgumentsFromFullOpWithScalarPass,
         ReplaceAtenAvgPoolWithJarvisAvgPoolPass,
         ReplaceWhereWithFullArgsWithWhereScalar,
-        ReplaceGeluWithApproximateGeluPass,
+        ReplaceAtenApproxGeluWithApproxGeluPass,
         ReplaceSplitWithSlicePass,
         ReplacePowWithMulPass,
     ]

backends/cadence/aot/tests/test_replace_ops_passes.py

@@ -32,7 +32,7 @@
     ReplaceConvWithIm2RowAndLinear,
     ReplaceEmptyTensorsWithFullPass,
     ReplaceFunctionallyEquivalentOpTargets,
-    ReplaceGeluWithApproximateGeluPass,
+    ReplaceAtenApproxGeluWithApproxGeluPass,
     ReplaceIm2RowWithViewPass,
     ReplaceLinearWithFullyConnectedOpPass,
     ReplaceMatmulWithTransposedMatmulPass,
@@ -1287,7 +1287,7 @@ def forward(self, cond: torch.Tensor):
             1,
         )
 
-    def test_replace_aten_gelu_with_approximate_gelu(self):
+    def test_no_replace_aten_gelu_with_approximate_gelu(self):
         class Gelu(torch.nn.Module):
             def forward(self, input):
                 return torch.nn.functional.gelu(input)
@@ -1296,7 +1296,28 @@ def forward(self, input):
 
         graph_module = export_to_edge(Gelu(), (inputs,)).exported_program().graph_module
 
-        p = ReplaceGeluWithApproximateGeluPass()
+        p = ReplaceAtenApproxGeluWithApproxGeluPass()
+        graph_after_passes = cast(PassResult, p(graph_module)).graph_module
+
+        # Assert that aten.gelu op was not decomposed, since it didn't have an approximate argument
+        self.assertEqual(
+            count_node(
+                graph_after_passes,
+                exir_ops.edge.aten.gelu.default,
+            ),
+            1,
+        )
+
+    def test_replace_aten_approximate_gelu_with_approximate_gelu(self):
+        class Gelu(torch.nn.Module):
+            def forward(self, input):
+                return torch.nn.functional.gelu(input, approximate = "tanh")
+
+        inputs = torch.randn(2, 1, 64)
+
+        graph_module = export_to_edge(Gelu(), (inputs,)).exported_program().graph_module
+
+        p = ReplaceAtenApproxGeluWithApproxGeluPass()
         graph_after_passes = cast(PassResult, p(graph_module)).graph_module
 
         # Assert that aten.gelu op was decomposed