Merge branch 'main' into an/openvino/quantizer_fix

Author: suryasidd

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

backends/arm/test/misc/test_debug_feats.py

@@ -16,6 +16,7 @@
 import torch
 from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
 from executorch.backends.arm.test import common
+from executorch.backends.arm.test.runner_utils import dbg_tosa_fb_to_json
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
     TosaPipelineFP,
@@ -238,25 +239,47 @@ def test_dump_tosa_debug_json(test_data: input_t1):
 
 @common.parametrize("test_data", Linear.inputs)
 def test_dump_tosa_debug_tosa(test_data: input_t1):
-    with tempfile.TemporaryDirectory() as tmpdir:
-        aten_ops: list[str] = []
-        exir_ops: list[str] = []
-        pipeline = TosaPipelineINT[input_t1](
-            module=Linear(),
-            test_data=test_data,
-            aten_op=aten_ops,
-            exir_op=exir_ops,
-            custom_path=tmpdir,
-            tosa_debug_mode=ArmCompileSpec.DebugMode.TOSA,
-        )
+    output_dir = "test_dump_tosa_debug"
 
-        pipeline.pop_stage("run_method_and_compare_outputs")
-        pipeline.run()
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
+    pipeline = TosaPipelineFP[input_t1](
+        module=Linear(),
+        test_data=test_data,
+        use_to_edge_transform_and_lower=True,
+        aten_op=aten_ops,
+        exir_op=exir_ops,
+        custom_path=output_dir,
+        tosa_debug_mode=ArmCompileSpec.DebugMode.TOSA,
+    )
 
-        json_output_path = Path(tmpdir) / "debug.json"
+    pipeline.pop_stage("run_method_and_compare_outputs")
+    pipeline.run()
+
+    output_path = Path(output_dir)
+    json_output_path = output_path / "debug.json"
+
+    # A JSON file should not be created when TOSA mode used
+    assert not json_output_path.exists()
+
+    # At least one TOSA file should exist
+    tosa_files = list(output_path.glob("*.tosa"))
+    assert len(tosa_files) > 0
+
+    tosa_file = tosa_files[0]
+    with tosa_file.open("rb") as f:
+        tosa_json = dbg_tosa_fb_to_json(f.read())
+
+    # Check all non-empty JSON strings are valid
+    ops = tosa_json["regions"][0]["blocks"][0]["operators"]
+    for op in ops:
+        if op["location"]["text"]:
+            try:
+                json.loads(op["location"]["text"])
+            except json.JSONDecodeError:
+                pytest.fail("Failed to load debug JSON string")
 
-        # A JSON file should not be created when TOSA mode used
-        assert not json_output_path.exists()
+    shutil.rmtree(output_dir, ignore_errors=True)
 
 
 @common.parametrize("test_data", Linear.inputs)

backends/arm/test/ops/test_clamp.py

@@ -35,6 +35,25 @@
     "rank_4_no_max": lambda: (torch.rand(1, 10, 10, 1) - 3, -3.3, None),
 }
 
+test_data_suite_int32 = {
+    "int32_rank2": lambda: (torch.randint(-50, 50, (2, 3), dtype=torch.int32), -10, 10),
+    "int32_rank3_no_min": lambda: (
+        torch.randint(-100, 100, (1, 3, 3), dtype=torch.int32),
+        None,
+        25,
+    ),
+    "int32_rank3_no_max": lambda: (
+        torch.randint(-100, 100, (1, 3, 3), dtype=torch.int32),
+        -25,
+        None,
+    ),
+    "int32_rank4_large_range": lambda: (
+        torch.randint(-200, 200, (1, 2, 4, 4), dtype=torch.int32),
+        torch.iinfo(torch.int32).min,
+        torch.iinfo(torch.int32).max,
+    ),
+}
+
 
 class Clamp(torch.nn.Module):
     def __init__(
@@ -53,7 +72,6 @@ def forward(self, x):
 
 @common.parametrize("test_data", test_data_suite)
 def test_clamp_tosa_FP(test_data):
-
     input_tensor, min_val, max_val = test_data()
     model = Clamp(min_val, max_val)
 
@@ -69,7 +87,6 @@ def test_clamp_tosa_FP(test_data):
 
 @common.parametrize("test_data", test_data_suite)
 def test_clamp_tosa_INT(test_data):
-
     input_tensor, min_val, max_val = test_data()
     model = Clamp(min_val, max_val)
 
@@ -84,6 +101,22 @@ def test_clamp_tosa_INT(test_data):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_suite_int32)
+def test_clamp_tosa_INT_int32_inputs(test_data):
+    input_tensor, min_val, max_val = test_data()
+    model = Clamp(min_val, max_val)
+
+    pipeline = TosaPipelineINT[input_t](
+        model,
+        (input_tensor,),
+        aten_op,
+        exir_op,
+    )
+    pipeline.change_args("run_method_and_compare_outputs", qtol=1)
+    pipeline.pop_stage("quantize")
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_suite)
 def test_clamp_tosa_INT_a16w8(test_data):
     """Test clamp operation with int16 I/O quantization for TOSA INT."""
@@ -103,7 +136,6 @@ def test_clamp_tosa_INT_a16w8(test_data):
 @common.parametrize("test_data", test_data_suite)
 @common.XfailIfNoCorstone300
 def test_clamp_u55_INT(test_data):
-
     input_tensor, min_val, max_val = test_data()
     model = Clamp(min_val, max_val)
 
@@ -140,7 +172,6 @@ def test_clamp_16a8w_u55_INT16(test_data):
 @common.parametrize("test_data", test_data_suite)
 @common.XfailIfNoCorstone320
 def test_clamp_u85_INT(test_data):
-
     input_tensor, min_val, max_val = test_data()
     model = Clamp(min_val, max_val)
 

backends/arm/test/passes/test_convert_int32_clamp_to_minmax_pass.py

@@ -0,0 +1,44 @@
+# 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 Tuple
+
+import torch
+from executorch.backends.arm._passes.decompose_int32_clamp_pass import (
+    DecomposeInt32ClampPass,
+)
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+
+input_t = Tuple[torch.Tensor]
+
+
+class ClampInt32(torch.nn.Module):
+    test_data = {"rand": (torch.randint(-50, 50, (2, 3), dtype=torch.int32),)}
+
+    def forward(self, x: torch.Tensor):
+        return torch.clamp(x, -10, 5)
+
+
+@common.parametrize("test_data", ClampInt32.test_data)
+def test_decompose_int32_clamp_pass(test_data: input_t):
+    module = ClampInt32()
+    pipeline = PassPipeline[input_t](
+        module,
+        test_data,
+        quantize=False,
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_clamp_default": 1,
+        },
+        ops_after_pass={
+            "executorch_exir_dialects_edge__ops_aten_minimum_default": 1,
+            "executorch_exir_dialects_edge__ops_aten_maximum_default": 1,
+        },
+        ops_not_after_pass=[
+            "executorch_exir_dialects_edge__ops_aten_clamp_default",
+        ],
+        pass_list=[DecomposeInt32ClampPass],
+    )
+    pipeline.run()

backends/arm/tosa/schemas/tosa_1.0.fbs

@@ -510,13 +510,31 @@ table TosaTensor {
   variable: bool;                   // is this a variable tensor
   is_unranked: bool;                // whether this is an unranked tensor
   variable_name:string;             // name for variable attribute
+
+  // In a model that is larger than 2GB, then tensors instead uses the following
+  // attributes to find stored data, which is outside of flatbuffers
+  // the offset is calculated relative to the beginning of the file and is only
+  // valid if > 1.
+  offset: ulong;
+  size: ulong;
+}
+
+table TosaShape {
+  name: string;                     // name of the shape
+  rank: uint32;                     // rank of the shape
+  data: [ubyte] (force_align: 8);   // raw data array if it's a constant shape
+}
+
+table OpLocation {
+  text: string;       // Opaque string, interpretted by user
 }
 
 table TosaOperator {
   op:Op;                    // operator enum
   attribute:Attribute;      // union structure. operator attribute
-  inputs:[string];          // list of input tensor names
-  outputs:[string];         // list of output tensor names
+  inputs:[string];          // list of input tensor or shape names
+  outputs:[string];         // list of output tensor or shape names
+  location: OpLocation;     // location of this Op in mlir
 }
 
 table TosaBasicBlock {
@@ -525,6 +543,7 @@ table TosaBasicBlock {
   tensors:[TosaTensor];     // tensors array
   inputs:[string];          // name of graph inputs
   outputs:[string];         // name of graph outputs
+  shapes:[TosaShape];       // shapes array
 }
 
 table TosaRegion {
@@ -537,4 +556,4 @@ table TosaGraph {
   regions:[TosaRegion];       // regions array
 }
 
-root_type TosaGraph;
+root_type TosaGraph;