Arm backend: Decompose integral float pow exponents (pytorch#19693)

Treat positive integral float scalar exponents like integer exponents in
DecomposeIntPowPass. This avoids lowering pow(x, 2.0) to TOSA POW, whose
reference model rejects negative bases even when the exponent is
mathematically integral.

Test Plan:
- lintrunner on changed files
- pytest test_decompose_int_pow_pass.py
- TinySwin2SR FP/INT TOSA reference smoke


cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Signed-off-by: Usamah Zaheer <usamah.zaheer@arm.com>

Author: usamahz

backends/arm/_passes/decompose_int_pow_pass.py

@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 
-from typing import Set, Type
+from typing import Optional, Set, Type
 
 from executorch.backends.arm._passes import ArmPass
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -21,6 +21,18 @@ class DecomposeIntPowPass(ArmPass):
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
+    @staticmethod
+    def _get_decomposable_integer_exponent(exp) -> Optional[int]:
+        if isinstance(exp, int):
+            return exp
+        # Exported models can represent positive integer-valued exponents as
+        # floats, for example pow(x, 2.0). Only exact values are decomposed:
+        # rounding near-integer floats would change fractional pow semantics,
+        # especially for negative bases.
+        if isinstance(exp, float) and exp > 0 and exp.is_integer():
+            return int(exp)
+        return None
+
     def call_operator(self, op, args, kwargs, meta):
         if op != exir_ops.edge.aten.pow.Tensor_Scalar:
             return super().call_operator(op, args, kwargs, meta)
@@ -43,9 +55,18 @@ def call_operator(self, op, args, kwargs, meta):
                 exir_ops.edge.aten.add.Tensor, (zeros, ones), {}, meta, True
             )
 
-        if not isinstance(exp, int):
+        exp = self._get_decomposable_integer_exponent(exp)
+        if exp is None:
             return super().call_operator(op, args, kwargs, meta)
 
+        if exp == 1:
+            ones = super().call_operator(
+                exir_ops.edge.aten.full_like.default, (x, 1), {}, meta, True
+            )
+            return super().call_operator(
+                exir_ops.edge.aten.mul.Tensor, (x, ones), {}, meta, True
+            )
+
         # Handle negative exponent
         if exp < 0:
             x = super().call_operator(

backends/arm/test/ops/test_pow.py

@@ -141,22 +141,11 @@ def test_pow_tensor_tensor_vgf_no_quant(test_data: Pow_TensorTensor.input_t):
     pipeline.run()
 
 
-x_fail = {
-    "exp_two": "TOSA constraints: If x <0 .",
-}
-
-x_fail_FP = {
-    "exp_two": "TOSA constraints: If x <0 .",
-}
-
-
 @common.parametrize(
     "test_data",
     Pow_TensorScalar.test_data
     | Pow_TensorScalar.test_data_fp16
     | Pow_TensorScalar.test_data_bf16,
-    xfails=x_fail_FP,
-    strict=False,
 )
 def test_pow_tensor_scalar_tosa_FP(test_data: Pow_TensorScalar.input_t):
     base, exp = test_data()
@@ -211,8 +200,6 @@ def test_pow_tensor_scalar_u85_INT(test_data: Pow_TensorScalar.input_t):
 @common.parametrize(
     "test_data",
     Pow_TensorScalar.test_data | Pow_TensorScalar.test_data_fp16,
-    x_fail_FP,
-    strict=False,
 )
 @common.SkipIfNoModelConverter
 def test_pow_tensor_scalar_vgf_no_quant(test_data: Pow_TensorScalar.input_t):

backends/arm/test/passes/test_decompose_int_pow_pass.py

@@ -35,7 +35,7 @@ def get_inputs(self) -> input_t:
 class Pow(torch.nn.Module):
     """Basic squaring."""
 
-    def __init__(self, exponent: int) -> None:
+    def __init__(self, exponent: int | float) -> None:
         super().__init__()
         self.exponent = exponent
 
@@ -48,12 +48,20 @@ def get_inputs(self) -> input_t:
 
 test_data: Dict[str, TestParam] = {
     "square": (Square(), 1),
+    "pow_1": (Pow(1), 1),
+    "pow_1_float": (Pow(1.0), 1),
     "pow_2": (Pow(2), 1),
+    "pow_2_float": (Pow(2.0), 1),
     "pow_3": (Pow(3), 2),
     "pow_0": (Pow(0), 0),
     "pow_neg_2": (Pow(-2), 1),
 }
 
+non_integer_float_test_data: Dict[str, ModuleWithInputs] = {
+    "pow_1_999999999": Pow(1.999999999),
+    "pow_2_000000001": Pow(2.000000001),
+}
+
 
 @common.parametrize("data", test_data)
 def test_decompose_int_pow_tosa_FP(data: TestParam) -> None:
@@ -74,3 +82,26 @@ def test_decompose_int_pow_tosa_FP(data: TestParam) -> None:
         pass_list=[DecomposeIntPowPass],
     )
     pipeline.run()
+
+
+@common.parametrize("module_with_inputs", non_integer_float_test_data)
+def test_decompose_int_pow_tosa_FP_non_integer_float(
+    module_with_inputs: ModuleWithInputs,
+) -> None:
+    module = cast(torch.nn.Module, module_with_inputs)
+    pow_op = "executorch_exir_dialects_edge__ops_aten_pow_Tensor_Scalar"
+    pipeline = PassPipeline[input_t](
+        module,
+        module_with_inputs.get_inputs(),
+        quantize=False,
+        ops_before_pass={
+            pow_op: 1,
+        },
+        ops_not_before_pass=[],
+        ops_after_pass={
+            pow_op: 1,
+            "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 0,
+        },
+        pass_list=[DecomposeIntPowPass],
+    )
+    pipeline.run()