Skip ReplacePowWithMulPass when pow promotes the input dtype (#24177)

backends/cadence/aot/replace_ops.py

@@ -2280,6 +2280,11 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
         x = node.args[0]
         assert isinstance(x, torch.fx.Node)
 
+        # mul keeps the input dtype, so it can't express a promoting pow
+        # (e.g. int32 ** 2.0 -> float32).
+        if x.meta["val"].dtype != node.meta["val"].dtype:
+            return False
+
         graph = node.graph
         result_node = x
 

backends/cadence/aot/tests/test_replace_ops_passes.py

@@ -326,6 +326,33 @@ def test_sub_replace_scalar_with_tensor_arg(
             0,
         )
 
+    @expand(
+        [
+            ["add", exir_ops.edge.aten.add.Scalar, exir_ops.edge.aten.add.Tensor],
+            ["sub", exir_ops.edge.aten.sub.Scalar, exir_ops.edge.aten.sub.Tensor],
+        ]
+    )
+    @torch.no_grad()
+    def test_replace_scalar_with_tensor_arg_positional_alpha(
+        self,
+        _,
+        scalar_op: torch._ops.OpOverload,
+        tensor_op: torch._ops.OpOverload,
+    ) -> None:
+        x = torch.randn(7, 5)
+        original_gm = single_op_builder(
+            placeholders=(x,),
+            op=scalar_op,
+            args=(x, 1.5, 2),
+        )
+        gm_before = copy.deepcopy(original_gm)
+        p = ReplaceScalarWithTensorArgPass()
+        graph_after_passes = cast(PassResult, p(original_gm)).graph_module
+
+        validate(gm_before, graph_after_passes, [x], "ReplaceScalarWithTensorArgPass")
+        self.assertEqual(count_node(graph_after_passes, tensor_op), 1)
+        self.assertEqual(count_node(graph_after_passes, scalar_op), 0)
+
     @expand(
         [
             ["3d", (7, 5, 6), 1.23],
@@ -2348,6 +2375,53 @@ def test_replace_pow_with_mul_not_applied(self, exponent: float) -> None:
             0,
         )
 
+    @expand([[2.0], [3.0]])
+    def test_replace_pow_with_mul_not_applied_on_type_promotion(
+        self, exponent: float
+    ) -> None:
+        # int32 ** float promotes to float32, but int32 * int32 stays int32.
+        x = torch.randint(0, 8, (2, 1, 64), dtype=torch.int32)
+        original_gm = single_op_builder(
+            placeholders=(x,),
+            op=exir_ops.edge.aten.pow.Tensor_Scalar,
+            args=(x, exponent),
+        )
+        p = ReplacePowWithMulPass()
+        result = cast(PassResult, p(original_gm))
+
+        self.assertFalse(result.modified)
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.pow.Tensor_Scalar),
+            1,
+        )
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.mul.Tensor), 0
+        )
+
+    @expand([[2], [3]])
+    def test_replace_pow_with_mul_int_input_int_exponent(self, exponent: int) -> None:
+        # int32 ** int stays int32, so the dtype guard still allows the rewrite.
+        x = torch.randint(0, 8, (2, 1, 64), dtype=torch.int32)
+        original_gm = single_op_builder(
+            placeholders=(x,),
+            op=exir_ops.edge.aten.pow.Tensor_Scalar,
+            args=(x, exponent),
+        )
+        gm_before = copy.deepcopy(original_gm)
+        p = ReplacePowWithMulPass()
+        result = cast(PassResult, p(original_gm))
+
+        self.assertTrue(result.modified)
+        validate(gm_before, result.graph_module, [x], "ReplacePowWithMulPass")
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.pow.Tensor_Scalar),
+            0,
+        )
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.mul.Tensor),
+            exponent - 1,
+        )
+
 
 class TestReplaceIm2rowWithViewPass(unittest.TestCase):
     def test_no_replacement_for_conv(self) -> None:

backends/transforms/replace_scalar_with_tensor.py

@@ -12,6 +12,7 @@
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.pass_base import ExportPass
 from torch._ops import OpOverload
+from torch.fx.operator_schemas import normalize_function
 
 
 class ReplaceScalarWithTensorArgPass(ExportPass):
@@ -71,8 +72,6 @@ def get_replacement(self, op, args, kwargs, meta):
                     },
                     meta=meta,
                 ),
-                # Other args.
-                *args[2:],
             ),
             kwargs=kwargs,
             meta=meta,
@@ -96,4 +95,10 @@ def call_operator(self, op, args, kwargs, meta):
         ):
             return super().call_operator(op, args, kwargs, meta)
 
-        return self.get_replacement(op, args, kwargs, meta)
+        # alpha is positional in add/sub.Scalar but keyword-only in .Tensor.
+        bound = normalize_function(
+            op, args, kwargs, normalize_to_only_use_kwargs=True
+        ).kwargs
+        return self.get_replacement(
+            op, (bound["input"], bound["other"]), {"alpha": bound["alpha"]}, meta
+        )