Implement a coversion pass: pow(2,x) to mul(x,x).

Differential Revision: D73405855

Pull Request resolved: #10373

Author: eigen-k

backends/cadence/aot/replace_ops.py

@@ -2259,6 +2259,34 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         return result
 
 
+
+@register_cadence_pass(CadencePassAttribute(opt_level=1))
+class ReplacePowWithMullPass(ExportPass):
+    """
+    Replace the pow op with degree 2 for a mul op.
+    """
+
+    def call_operator(
+        self,
+        op,
+        args: Tuple[Argument, ...],
+        kwargs: Dict[str, Argument],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        # TODO(eigen): Add support for other degrees.
+        if op not in {
+            exir_ops.edge.aten.pow.Scalar,
+        } or args[0] != 2:
+            return super().call_operator(op, args, kwargs, meta)
+
+        return super().call_operator(
+            exir_ops.edge.aten.mul.Tensor,
+            (args[1], args[1]),
+            {},
+            meta,
+        )
+
+
 # This class encapsulates all the functions that replace/switch one op in the
 # graph with another.
 class CadenceReplaceOpsInGraph:
@@ -2299,4 +2327,5 @@ class CadenceReplaceOpsInGraph:
         ReplaceWhereWithFullArgsWithWhereScalar,
         ReplaceGeluWithApproximateGeluPass,
         ReplaceSplitWithSlicePass,
+        ReplacePowWithMullPass,
     ]

backends/cadence/aot/tests/test_replace_ops_passes.py

@@ -30,6 +30,7 @@
     ReplaceEmptyTensorsWithFullPass,
     ReplaceFunctionallyEquivalentOpTargets,
     ReplaceGeluWithApproximateGeluPass,
+    ReplacePowWithMullPass,
     ReplaceIm2RowWithViewPass,
     ReplaceLinearWithFullyConnectedOpPass,
     ReplaceMMWithAddMMPass,
@@ -1334,6 +1335,35 @@ def test_replace_split_with_sizes_with_slice(self):
             2,
         )
 
+    def test_replace_pow_with_mul(self):
+        class Pow(torch.nn.Module):
+            def forward(self, input):
+                return  torch.ops.aten.pow.Scalar(2, input)
+
+        input = torch.randn(2, 1, 64)
+
+        graph_module = export_to_edge(Pow(), (input,)).exported_program().graph_module
+
+        p = ReplacePowWithMullPass()
+        graph_after_passes = cast(PassResult, p(graph_module)).graph_module
+
+
+        self.assertEqual(
+            count_node(
+                graph_after_passes,
+                exir_ops.edge.aten.pow.Scalar,
+            ),
+            0,
+        )
+
+        self.assertEqual(
+            count_node(
+                graph_after_passes,
+                exir_ops.edge.aten.mul.Tensor,
+            ),
+            1,
+        )
+
 
 class TestReplaceIm2rowWithViewPass(unittest.TestCase):
     def test_no_replacement_for_conv(self):