Merge branch 'main' into pr_model_improve

Author: zingo

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/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/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/tosa_supported_operators.py

@@ -176,8 +176,10 @@ def is_node_supported(
             exir_ops.edge.aten.full_like.default,
             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.mul.Tensor,
             exir_ops.edge.aten.add.Scalar,
             exir_ops.edge.aten.sub.Scalar,
@@ -194,6 +196,7 @@ def is_node_supported(
             exir_ops.edge.aten.reciprocal.default,
             exir_ops.edge.aten.relu.default,
             exir_ops.edge.aten.leaky_relu.default,
+            exir_ops.edge.aten.sqrt.default,
             exir_ops.edge.aten.rsqrt.default,
             exir_ops.edge.aten._softmax.default,
             exir_ops.edge.aten.select_copy.int,
@@ -256,6 +259,7 @@ def is_node_supported(
                 exir_ops.edge.aten.var.correction,
                 exir_ops.edge.aten.var.dim,
                 exir_ops.edge.aten.add.Scalar,
+                exir_ops.edge.aten.sqrt.default,
                 exir_ops.edge.aten.sub.Scalar,
                 exir_ops.edge.aten.mul.Scalar,
                 exir_ops.edge.aten.div.Scalar,

backends/arm/operators/op_clamp.py

@@ -63,7 +63,8 @@ def cast_type(value: Any) -> int | float:
                 # Attempt to cast to float
                 return float(value)
 
-        assert 2 <= len(node.args) <= 3
+        if len(node.args) != 2 and len(node.args) != 3:
+            raise ValueError(f"Expected len(node.args) to be 2 or 3, got {node.args}")
 
         min_arg = dtype_min
         max_arg = dtype_max
@@ -84,7 +85,10 @@ def define_node(
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
-        assert len(node.all_input_nodes) == 1
+        if len(node.all_input_nodes) != 1:
+            raise ValueError(
+                f"Expected 1 input for {self.target}, got {len(node.all_input_nodes)}"
+            )
 
         min_int8, max_int8 = self._get_min_max_arguments(
             node,
@@ -122,7 +126,10 @@ def define_node(
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
-        assert len(node.all_input_nodes) == 1
+        if len(node.all_input_nodes) != 1:
+            raise ValueError(
+                f"Expected 1 input for {self.target}, got {len(node.all_input_nodes)}"
+            )
 
         if inputs[0].dtype == ts.DType.INT8:
             # Call the inherited define_node for handling integers

backends/arm/operators/op_minimum.py

@@ -37,20 +37,27 @@ def define_node(
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
-        assert inputs[0].dtype == inputs[1].dtype
+        if inputs[0].dtype != inputs[1].dtype and inputs[0].dtype != output.dtype:
+            raise TypeError(
+                f"Data type of inputs and output must be the same. Got input 0 dtype: "
+                f"{inputs[0].dtype}, input 1 dtype: {inputs[1].dtype} and output "
+                f"dtype: {output.dtype}"
+            )
 
         scale_back = 1.0
         min_output = output
         if inputs[0].dtype == ts.DType.INT8:
             input_qparams = get_input_qparams(node)
-            assert (
-                len(input_qparams) == 2
-            ), f"Both inputs needs to have quantization information for {node}"
-            # insert RESCALEs to int32
-            assert (
-                input_qparams[0] == input_qparams[1]
-            ), "Both inputs must have same quantization for MIN"
+            if len(input_qparams) != 2:
+                raise ValueError(
+                    f"Both inputs need to have quantization information for {node}"
+                )
+            if input_qparams[0] != input_qparams[1]:
+                raise ValueError(
+                    "Both inputs must have the same quantization parameters for MIN"
+                )
 
+            # insert RESCALEs to int32
             operand_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(
                 tosa_graph, inputs, node
             )

backends/arm/process_node.py

@@ -15,6 +15,14 @@
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.backends.arm.tosa_utils import getNodeArgs, tosa_shape
+from torch._export.utils import (
+    get_buffer,
+    get_lifted_tensor_constant,
+    get_param,
+    is_buffer,
+    is_lifted_tensor_constant,
+    is_param,
+)
 from torch.export.exported_program import ExportedProgram
 
 
@@ -99,8 +107,7 @@ def process_inputs_to_parameters(
             f"Failed processing parameter placeholder: {node.name}. "
             "Is the original torch function supported?"
         ) from e
-    parameter_name = edge_program.graph_signature.inputs_to_parameters[tosa_arg.name]
-    parameter_data = edge_program.state_dict[parameter_name]
+    parameter_data = get_param(edge_program, node)
 
     assert isinstance(parameter_data, torch.Tensor), "Expect Attr to be tensor"
     parameter_values = parameter_data.detach().numpy()
@@ -128,8 +135,7 @@ def process_inputs_to_buffers(
             f"Failed processing buffer placeholder: {node.name}. "
             "Is the original torch function supported?"
         ) from e
-    buffer_name = edge_program.graph_signature.inputs_to_buffers[node.name]
-    buffer_data = edge_program.state_dict[buffer_name]
+    buffer_data = get_buffer(edge_program, node)
 
     assert isinstance(buffer_data, torch.Tensor), "Expect Attr to be tensor"
     buffer_values = buffer_data.detach().numpy()
@@ -156,11 +162,8 @@ def process_inputs_to_lifted_tensor_constants(
             f"Failed processing lifted tensor constant placeholder: {node.name}. "
             "Is the original torch function supported?"
         ) from e
-    tensor_name = edge_program.graph_signature.inputs_to_lifted_tensor_constants[
-        tosa_arg.name
-    ]
-    tensor = edge_program.tensor_constants[tensor_name]
-    tensor_data = tensor.detach().numpy()
+    tensor = get_lifted_tensor_constant(edge_program, node)
+    tensor_data = tensor.detach().numpy()  # type: ignore[union-attr]
 
     tosa_graph.addConst(
         tensor_data.shape, tosa_arg.dtype, tensor_data, name=tosa_arg.name
@@ -179,11 +182,11 @@ def process_placeholder(
 
     if node.name in edge_program.graph_signature.user_inputs:
         process_inputs(node, tosa_graph, tosa_spec)
-    elif node.name in edge_program.graph_signature.inputs_to_parameters:
+    elif is_param(edge_program, node):
         process_inputs_to_parameters(node, tosa_graph, edge_program, tosa_spec)
-    elif node.name in edge_program.graph_signature.inputs_to_buffers:
+    elif is_buffer(edge_program, node):
         process_inputs_to_buffers(node, tosa_graph, edge_program)
-    elif node.name in edge_program.graph_signature.inputs_to_lifted_tensor_constants:
+    elif is_lifted_tensor_constant(edge_program, node):
         process_inputs_to_lifted_tensor_constants(node, tosa_graph, edge_program)
     elif node.name in edge_program.graph_signature.inputs_to_lifted_custom_objs:
         raise NotImplementedError(

backends/arm/test/misc/test_non_persistent_buffers.py

@@ -0,0 +1,49 @@
+# 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.
+
+import torch
+import torch.nn as nn
+
+from executorch.backends.arm.test.common import parametrize
+from executorch.backends.arm.test.tester.test_pipeline import (
+    TosaPipelineBI,
+    TosaPipelineMI,
+)
+
+
+class NonPersistentBuffer(nn.Module):
+    """
+    Min code version registering a non-persistent input buffer.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("test_buff", torch.rand(2, 2, 2, 2), persistent=False)
+
+    def forward(self, x):
+        return x - self.test_buff
+
+
+test_input = {"input": (torch.ones(2, 2, 2, 2),)}
+
+input_t = tuple[torch.Tensor]
+
+
+@parametrize("test_data", test_input)
+def test_non_persistent_buffer_MI(test_data: input_t):
+    """
+    Test validates Arm backend handling of non-persistent buffers
+    and ensures that there are no asserts or errors when they are used.
+    """
+    TosaPipelineMI[input_t](NonPersistentBuffer(), test_data, "").run()
+
+
+@parametrize("test_data", test_input)
+def test_non_persistent_buffer_BI(test_data: input_t):
+    """
+    Test validates Arm backend handling of non-persistent buffers
+    and ensures that there are no asserts or errors when they are used.
+    """
+    TosaPipelineBI[input_t](NonPersistentBuffer(), test_data, "").run()

backends/arm/test/models/test_llama.py

@@ -79,24 +79,6 @@ def prepare_model(self):
 
         llama_model, llama_inputs, llama_meta = get_llama_model(args)
 
-        # TODO: Remove workaround since attention mask should not be persistent,
-        # it only works if input shape is always the same
-        freqs_c = "freqs_cos"
-        freqs_s = "freqs_sin"
-        for i in range(llama_model.n_layers):
-            val = llama_model.layers[i].attention.get_buffer("mask")
-            llama_model.layers[i].attention.register_buffer(
-                "mask", val, persistent=True
-            )
-            val = llama_model.layers[i].attention.rope.get_buffer(freqs_c)
-            llama_model.layers[i].attention.rope.register_buffer(
-                freqs_c, val, persistent=True
-            )
-            val = llama_model.layers[i].attention.rope.get_buffer(freqs_s)
-            llama_model.layers[i].attention.rope.register_buffer(
-                freqs_s, val, persistent=True
-            )
-
         return llama_model, llama_inputs, llama_meta
 
     def test_llama_tosa_MI(self):