diff --git a/backends/cadence/aot/replace_ops.py b/backends/cadence/aot/replace_ops.py index f68b05857d2..bc04f749491 100644 --- a/backends/cadence/aot/replace_ops.py +++ b/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 diff --git a/backends/cadence/aot/tests/test_replace_ops_passes.py b/backends/cadence/aot/tests/test_replace_ops_passes.py index 6d1e34060ff..0b6c07bddee 100644 --- a/backends/cadence/aot/tests/test_replace_ops_passes.py +++ b/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: diff --git a/backends/transforms/replace_scalar_with_tensor.py b/backends/transforms/replace_scalar_with_tensor.py index ed45998dc56..3ff83e0e3bf 100644 --- a/backends/transforms/replace_scalar_with_tensor.py +++ b/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 + )