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/cadence/hifi/operators/op_div.cpp b/backends/cadence/hifi/operators/op_div.cpp index 76efce98c3b..9923a99db96 100644 --- a/backends/cadence/hifi/operators/op_div.cpp +++ b/backends/cadence/hifi/operators/op_div.cpp @@ -210,6 +210,10 @@ Tensor& div_out_mode( Tensor& out) { ET_DCHECK_MSG(!tensor_has_zero(b), "divisor tensor contains zero"); + if (!mode.has_value()) { + return div_out(ctx, a, b, out); + } + ET_KERNEL_CHECK( ctx, torch::executor::resize_to_broadcast_target_size(a, b, out) == Error::Ok, diff --git a/backends/cadence/hifi/operators/tests/test_op_div.cpp b/backends/cadence/hifi/operators/tests/test_op_div.cpp index 7237e0aeb24..955b9c6bfc0 100644 --- a/backends/cadence/hifi/operators/tests/test_op_div.cpp +++ b/backends/cadence/hifi/operators/tests/test_op_div.cpp @@ -59,6 +59,31 @@ TEST_F(HiFiDivTest, DISABLED_Int32FloorDivideTest) { EXPECT_TENSOR_EQ(out, tf.full(sizes, kExpectedResult)); } +TEST_F(HiFiDivTest, FloatNoRoundingModeIsTrueDivision) { + TensorFactory tf; + const std::vector sizes{2, 2}; + Tensor out = tf.zeros(sizes); + const Tensor numerator = tf.make(sizes, {7.0, -7.0, 1.0, -1.0}); + const Tensor denominator = tf.make(sizes, {2.0, 2.0, 4.0, 4.0}); + + div_out_mode(numerator, denominator, std::nullopt, out); + + EXPECT_TENSOR_EQ(out, tf.make(sizes, {3.5, -3.5, 0.25, -0.25})); +} + +TEST_F(HiFiDivTest, Int32NoRoundingModeIsTrueDivisionToFloat) { + TensorFactory tf_int; + TensorFactory tf_float; + const std::vector sizes{2, 2}; + Tensor out = tf_float.zeros(sizes); + const Tensor numerator = tf_int.make(sizes, {7, -7, 1, -1}); + const Tensor denominator = tf_int.make(sizes, {2, 2, 4, 4}); + + div_out_mode(numerator, denominator, std::nullopt, out); + + EXPECT_TENSOR_EQ(out, tf_float.make(sizes, {3.5, -3.5, 0.25, -0.25})); +} + } // namespace } // namespace native } // namespace HiFi diff --git a/backends/cadence/utils/facto_util.py b/backends/cadence/utils/facto_util.py index 6ac884e9603..8ae2ad6a722 100644 --- a/backends/cadence/utils/facto_util.py +++ b/backends/cadence/utils/facto_util.py @@ -518,6 +518,17 @@ def facto_testcase_gen( # noqa: C901 spec.inspec[index].constraints.extend( [cp.Value.Ge(lambda deps, _: deps[1])] ) + elif in_spec.name == "exponent" and op_name == "pow.Tensor_Scalar": + # A Gt(0) bound makes nextup(0) = 5e-324 an edge value, which + # flushes to 0 in float32 kernels and turns 0 ** e into 0 ** 0. + spec.inspec[index].constraints.extend( + [ + cp.Value.Ge( + lambda deps, dtype: torch.finfo(torch.float32).tiny + ), + cp.Value.Le(lambda deps, dtype: 2), + ] + ) elif in_spec.name == "negative_slope" and op_name == "leaky_relu.default": # For leaky_relu, negative_slope should be in typical range (0, 1] spec.inspec[index].constraints.extend( 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 + )