Keep the FACTO pow exponent representable in float32 (#24179)

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/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,

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<ScalarType::Float> tf;
+  const std::vector<int> 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<ScalarType::Int> tf_int;
+  TensorFactory<ScalarType::Float> tf_float;
+  const std::vector<int> 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

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(

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
+        )