Merge branch 'main' into experimental/clone_support

Author: Sebastian-Larsson

backends/arm/_passes/__init__.py

@@ -52,6 +52,7 @@
 from .decompose_int16_activation_conv2d_pass import (  # noqa
     DecomposeConv2dWithInt16ActivationPass,
 )
+from .decompose_int32_clamp_pass import DecomposeInt32ClampPass  # noqa
 from .decompose_int_pow_pass import DecomposeIntPowPass  # noqa
 from .decompose_layernorm_pass import DecomposeLayerNormPass  # noqa
 from .decompose_leaky_relu_pass import DecomposeLeakyReLUPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -55,6 +55,7 @@
     DecomposeGluPass,
     DecomposeGroupedConvPass,
     DecomposeGroupNormPass,
+    DecomposeInt32ClampPass,
     DecomposeIntPowPass,
     DecomposeLayerNormPass,
     DecomposeLeakyReLUPass,
@@ -122,7 +123,6 @@
 
 
 class ArmPassManager(PassManager):
-
     def __init__(self, tosa_spec: TosaSpecification) -> None:
         self.tosa_spec = tosa_spec
         super().__init__()
@@ -174,6 +174,7 @@ def _tosa_pipeline(
                 FuseQuantizedActivationPass(),
                 RemoveGetItemPass(),
                 ConvertToClampPass(),
+                DecomposeInt32ClampPass(),
                 DecomposeGroupNormPass(),
                 DecomposeLayerNormPass(),
                 DecomposeBatchNormNoStatsPass(),

backends/arm/_passes/decompose_int32_clamp_pass.py

@@ -0,0 +1,72 @@
+# 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.
+
+from typing import Set, Type
+
+import torch
+from executorch.backends.arm._passes import ArmPass
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+
+class DecomposeInt32ClampPass(ArmPass):
+    """Rewrite int32 clamp into min/max chain since TOSA lacks int32 clamp support."""
+
+    _passes_required_after: Set[Type[ExportPass]] = set()
+    _supported_ops = {
+        exir_ops.edge.aten.clamp.default,
+        torch.ops.aten.clamp.default,
+    }
+
+    def _ensure_tensor(
+        self,
+        value,
+        ref_tensor,
+        dtype,
+        rank,
+        meta,
+    ):
+        if value is None:
+            return None
+        return super().call_operator(
+            exir_ops.edge.aten.full.default,
+            ((1,) * rank, value),
+            {"dtype": dtype},
+            meta,
+            updated=True,
+        )
+
+    def call_operator(self, op, args, kwargs, meta):
+        val = meta["val"]
+        if op not in self._supported_ops or val.dtype != torch.int32:
+            return super().call_operator(op, args, kwargs, meta)
+
+        input_tensor = args[0]
+        min_arg = args[1] if len(args) > 1 else None
+        max_arg = args[2] if len(args) > 2 else None
+        dtype = val.dtype
+        rank = len(val.shape)
+
+        min_arg = self._ensure_tensor(min_arg, input_tensor, dtype, rank, meta)
+        max_arg = self._ensure_tensor(max_arg, input_tensor, dtype, rank, meta)
+
+        current = input_tensor
+        if max_arg is not None:
+            current = super().call_operator(
+                exir_ops.edge.aten.minimum.default,
+                (current, max_arg),
+                {},
+                meta,
+                updated=True,
+            )
+        if min_arg is not None:
+            current = super().call_operator(
+                exir_ops.edge.aten.maximum.default,
+                (current, min_arg),
+                {},
+                meta,
+                updated=True,
+            )
+        return current

backends/arm/operators/node_visitor.py

@@ -87,7 +87,7 @@ def _serialize_operator(
             None: Mutates ``tosa_graph`` in place.
 
         """
-        op_location = ts.TosaOpLocation()
+        op_location = None
         if self.debug_hook:
             debug_info = self.debug_hook.add(
                 node,
@@ -96,7 +96,7 @@ def _serialize_operator(
             )
 
             if self.debug_hook.mode == ArmCompileSpec.DebugMode.TOSA:
-                op_location.text = json.dumps(debug_info.to_dict())
+                op_location = json.dumps(debug_info.to_dict())
 
         tosa_graph.addOperator(
             tosa_op,

backends/arm/operators/op_clamp.py

@@ -40,7 +40,6 @@ def __init__(self, *args):
     def _get_min_max_arguments(
         self, node: Node, dtype: torch.dtype
     ) -> Tuple[int | float, int | float]:
-
         def cast_type(value: Any) -> int | float:
             if isinstance(value, int):
                 return value
@@ -91,7 +90,12 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.FP16, ts.DType.FP32],
+            [
+                ts.DType.INT8,
+                ts.DType.INT16,
+                ts.DType.FP16,
+                ts.DType.FP32,
+            ],
             output.tosa_spec,
         )