Merge branch 'main' into docs/jhelsby/new-contributor-guide-update

Author: jhelsby

backends/arm/_passes/__init__.py

@@ -27,6 +27,7 @@
 from .decompose_select import DecomposeSelectPass  # noqa
 from .decompose_softmax_pass import DecomposeSoftmaxPass  # noqa
 from .decompose_softmax_unstable_pass import DecomposeSoftmaxUnstablePass  # noqa
+from .decompose_sqrt_pass import DecomposeSqrtPass  # noqa
 from .decompose_var_pass import DecomposeVarPass  # noqa
 from .fold_qdq_with_annotated_qparams_pass import (  # noqa
     FoldAndAnnotateQParamsPass,

backends/arm/_passes/arm_pass_manager.py

@@ -32,6 +32,7 @@
     DecomposeSelectPass,
     DecomposeSoftmaxPass,
     DecomposeSoftmaxUnstablePass,
+    DecomposeSqrtPass,
     DecomposeVarPass,
     FoldAndAnnotateQParamsPass,
     FuseBatchnorm2DPass,
@@ -115,6 +116,7 @@ def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         return self._transform(exported_program.graph_module)
 
     def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
+        self.add_pass(DecomposeSqrtPass())
         self.add_pass(ReplaceScalarWithTensorArgPassTOSAMI())
         self.add_pass(FuseQuantizedActivationPass())
         self.add_pass(RemoveGetItemPass())
@@ -181,6 +183,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         self.add_pass(DecomposeMeanDimPass())
         self.add_pass(DecomposeDivPass())
         self.add_pass(DecomposeLeakyReLUPass())
+        self.add_pass(DecomposeSqrtPass())
 
         if isinstance(self.tosa_spec, Tosa_0_80) and self.tosa_spec.is_U55_subset:
             # Numerically stable softmax uses amax which is not supported on Ethos-U55

backends/arm/_passes/cast_int64_pass.py

@@ -12,7 +12,6 @@
 from torch._export.utils import is_buffer
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.WARNING)
 
 
 class CastInt64BuffersToInt32Pass(ExportPass):

backends/arm/_passes/decompose_sqrt_pass.py

@@ -0,0 +1,39 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+import torch
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+edge_sqrt_ops = (exir_ops.edge.aten.sqrt.default,)
+aten_sqrt_ops = (
+    torch.ops.aten.sqrt.default,
+    torch.ops.aten.sqrt_.default,
+)
+
+
+def get_sqrt_decomposition(op) -> tuple:
+    # TODO : "MLETORCH-863 : Replace current sqrt -> pow.Tensor_Scalar workaround with pow.Tensor_Tensor"
+    if op in edge_sqrt_ops:
+        return exir_ops.edge.aten.pow.Tensor_Scalar
+    if op in aten_sqrt_ops:
+        return torch.ops.aten.pow.Tensor_Scalar
+    raise RuntimeError(f"Can't get sqrt decomposition for op {op}")
+
+
+class DecomposeSqrtPass(ExportPass):
+
+    def call_operator(self, op, args, kwargs, meta):
+        """
+        Decomposes `sqrt(x)` into `pow(x, 0.5)` for backend support.
+        """
+
+        if op not in (edge_sqrt_ops + aten_sqrt_ops):
+            return super().call_operator(op, args, kwargs, meta)
+
+        pow_op = get_sqrt_decomposition(op)
+
+        return super().call_operator(pow_op, (args[0], 0.5), {}, meta)

backends/arm/_passes/match_arg_ranks_pass.py

@@ -48,6 +48,8 @@ def __init__(self, exported_program):
         exir_ops.edge.aten.bitwise_right_shift.Tensor,
         exir_ops.edge.aten.bitwise_left_shift.Tensor,
         exir_ops.edge.aten.eq.Tensor,
+        exir_ops.edge.aten.gt.Tensor,
+        exir_ops.edge.aten.lt.Tensor,
         exir_ops.edge.aten.pow.Tensor_Tensor,
         exir_ops.edge.aten.where.self,
     ]

backends/arm/_passes/replace_scalar_with_tensor_pass.py

@@ -26,13 +26,17 @@
     exir_ops.edge.aten.__rshift__.Scalar: exir_ops.edge.aten.bitwise_right_shift.Tensor,
     exir_ops.edge.aten.__lshift__.Scalar: exir_ops.edge.aten.bitwise_left_shift.Tensor,
     exir_ops.edge.aten.eq.Scalar: exir_ops.edge.aten.eq.Tensor,
+    exir_ops.edge.aten.gt.Scalar: exir_ops.edge.aten.gt.Tensor,
+    exir_ops.edge.aten.lt.Scalar: exir_ops.edge.aten.lt.Tensor,
     torch.ops.aten.add.Scalar: torch.ops.aten.add.Tensor,
     torch.ops.aten.sub.Scalar: torch.ops.aten.sub.Tensor,
     torch.ops.aten.mul.Scalar: torch.ops.aten.mul.Tensor,
     torch.ops.aten.div.Scalar: torch.ops.aten.div.Tensor,
     torch.ops.aten.__rshift__.Scalar: torch.ops.aten.bitwise_right_shift.Tensor,
     torch.ops.aten.__lshift__.Scalar: torch.ops.aten.bitwise_left_shift.Tensor,
     torch.ops.aten.eq.Scalar: torch.ops.aten.eq.Tensor,
+    torch.ops.aten.gt.Scalar: torch.ops.aten.gt.Tensor,
+    torch.ops.aten.lt.Scalar: torch.ops.aten.lt.Tensor,
 }
 
 

backends/arm/arm_backend.py

@@ -11,19 +11,13 @@
 # JIT compiler flows.
 #
 
-import logging
-
 from typing import List, Optional
 
 from executorch.backends.arm.tosa_specification import TosaSpecification
 
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.WARNING)
-
-
 class ArmCompileSpecBuilder:
     def __init__(self):
         self.compile_spec: List[CompileSpec] = []

backends/arm/operator_support/ethos_u55_support.py

@@ -135,8 +135,10 @@ class EthosU55NotSupported(OperatorSupportBase):
         exir_ops.edge.aten.eq.Scalar,
         exir_ops.edge.aten.ge.Tensor,
         exir_ops.edge.aten.gt.Tensor,
+        exir_ops.edge.aten.gt.Scalar,
         exir_ops.edge.aten.le.Tensor,
         exir_ops.edge.aten.lt.Tensor,
+        exir_ops.edge.aten.lt.Scalar,
         exir_ops.edge.aten.flip.default,  # REVERSE
         exir_ops.edge.aten.grid_sampler_2d,  # GATHER
         exir_ops.edge.aten.scatter.src,

backends/arm/operator_support/right_shift_support.py

@@ -17,7 +17,6 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.WARNING)
 
 
 @register_tosa_support_check

backends/arm/operator_support/slice_copy_support.py

@@ -16,7 +16,6 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.WARNING)
 
 
 @register_tosa_support_check