Arm backend: Update composable_quantizer to latest changes (#20912)

Updates the composable_quantizer to include corresponding changes as in
the original one to allow a simple transition.

- Adds constant ops support
- Adds moveaxis/movedims support
- Adds special to handling support
- Update conv-relu annotation logic to match old behavior

Signed-off-by: Adrian Lundell <adrian.lundell@arm.com>

Author: AdrianLundell

backends/arm/quantizer/quantization_config.py

@@ -21,6 +21,7 @@
 
 from torchao.quantization.pt2e.quantizer import (
     DerivedQuantizationSpec,
+    FixedQParamsQuantizationSpec,
     QuantizationSpec,
     QuantizationSpecBase,
     SharedQuantizationSpec,
@@ -355,6 +356,7 @@ def get_output_act_qspec(
 
         If node is a pooling or upsample operator, returns a shared quantization spec.
         If no weight spec is configured, return ``None``.
+        If node is a `to.dtype` operator, returns a fixed quantization spec if the input is integer and the output is floating-point.
 
         """
 
@@ -373,6 +375,39 @@ def get_output_act_qspec(
             return _get_fixed_qparams_qspec(
                 node.target, _fixed_output_qspec_ops, output_act_qspec
             )
+        if node.target == torch.ops.aten.to.dtype:
+            from executorch.backends.arm.quantizer.quantizer_support import CastCheck
+
+            input_node = node.all_input_nodes[0]
+            input_val = input_node.meta.get("val", None)
+            output_val = node.meta.get("val", None)
+            if (
+                isinstance(input_val, torch.Tensor)
+                and isinstance(output_val, torch.Tensor)
+                and CastCheck.is_integer_to_integer(input_val.dtype, output_val.dtype)
+            ):
+                return None
+            if (
+                isinstance(input_val, torch.Tensor)
+                and isinstance(output_val, torch.Tensor)
+                and CastCheck.is_integer_to_float(input_val.dtype, output_val.dtype)
+            ):
+                return FixedQParamsQuantizationSpec(
+                    dtype=input_val.dtype,
+                    scale=1.0,
+                    zero_point=0,
+                    quant_min=torch.iinfo(input_val.dtype).min,
+                    quant_max=torch.iinfo(input_val.dtype).max,
+                    qscheme=torch.per_tensor_symmetric,
+                    is_dynamic=False,
+                )
+            if (
+                isinstance(input_val, torch.Tensor)
+                and isinstance(output_val, torch.Tensor)
+                and CastCheck.is_float_identity(input_val.dtype, output_val.dtype)
+            ):
+                return SharedQuantizationSpec((input_node, node))
+            return super().get_output_act_qspec(node)
         if node.target not in self.SHARED_OUTPUT_ACT_QSPEC_PATTERNS:
             return super().get_output_act_qspec()
         if len(node.args) == 0:

backends/arm/quantizer/quantizer_support.py

@@ -21,6 +21,19 @@ def combo_pattern(*pattern_lists):
 
 
 class ReluFusedPatternCheck(PatternCheck):
+    @classmethod
+    def check_pattern(cls, pattern):
+        output_node = pattern[-1] if pattern else None
+        if output_node is None:
+            return False
+
+        min_val = float(output_node.args[1]) if len(output_node.args) > 1 else -1.0
+        return (
+            output_node.target
+            not in (torch.ops.aten.hardtanh.default, torch.ops.aten.hardtanh_.default)
+            or min_val == 0
+        )
+
     @classmethod
     def check_quantization_config(cls, pattern, quantization_config):
         if quantization_config is None:
@@ -55,6 +68,47 @@ def check_pattern(cls, pattern):
         return True
 
 
+class CastCheck(PatternCheck):
+    INTEGER_CAST_DTYPES = (torch.int8, torch.int16, torch.int32)
+
+    @classmethod
+    def is_integer_to_integer(cls, input_dtype: torch.dtype, output_dtype: torch.dtype):
+        return (
+            input_dtype in cls.INTEGER_CAST_DTYPES
+            and output_dtype in cls.INTEGER_CAST_DTYPES
+        )
+
+    @classmethod
+    def is_integer_to_float(cls, input_dtype: torch.dtype, output_dtype: torch.dtype):
+        return input_dtype in cls.INTEGER_CAST_DTYPES and output_dtype.is_floating_point
+
+    @classmethod
+    def is_float_identity(cls, input_dtype: torch.dtype, output_dtype: torch.dtype):
+        return input_dtype.is_floating_point and input_dtype == output_dtype
+
+    @classmethod
+    def is_supported_cast(cls, input_dtype: torch.dtype, output_dtype: torch.dtype):
+        return (
+            cls.is_integer_to_integer(input_dtype, output_dtype)
+            or cls.is_integer_to_float(input_dtype, output_dtype)
+            or cls.is_float_identity(input_dtype, output_dtype)
+        )
+
+    @classmethod
+    def check_pattern(cls, pattern):
+        node = pattern[0]
+        if len(node.all_input_nodes) == 0:
+            return False
+
+        try:
+            input_tensor = get_first_fake_tensor(node.all_input_nodes[0])
+            output_tensor = get_first_fake_tensor(node)
+        except Exception:
+            return False
+
+        return cls.is_supported_cast(input_tensor.dtype, output_tensor.dtype)
+
+
 BINARY_OP_PATTERNS = [
     (torch.ops.aten.add.Tensor,),
     (torch.ops.aten.add_.Tensor,),
@@ -77,8 +131,6 @@ def check_pattern(cls, pattern):
     torch.ops.aten.relu_.default,
     torch.ops.aten.hardtanh.default,
     torch.ops.aten.hardtanh_.default,
-    torch.ops.aten.clamp.default,
-    torch.ops.aten.clamp_.default,
 ]
 BATCH_NORM_OPS = [torch.ops.aten.batch_norm.default]
 LINEAR_OP_PATTERNS = (
@@ -175,6 +227,13 @@ def check_pattern(cls, pattern):
         (torch.ops.aten.atanh.default,),
         (torch.ops.aten.einsum.default,),
         (torch.ops.aten.grid_sampler.default,),
+        (torch.ops.aten.linspace.default,),
+        (torch.ops.aten.eye.default,),
+        (torch.ops.aten.moveaxis.int,),
+        (torch.ops.aten.moveaxis.intlist,),
+        (torch.ops.aten.movedim.int,),
+        (torch.ops.aten.movedim.intlist,),
+        (torch.ops.aten.to.dtype,),
     ]
 )
 TOSA_QUANTIZER_SUPPORT_DICT: dict[tuple[OpOverload, ...], type[PatternCheck] | None] = {
@@ -184,3 +243,4 @@ def check_pattern(cls, pattern):
     TOSA_QUANTIZER_SUPPORT_DICT[pattern] = ReluFusedPatternCheck
 for pattern in BINARY_OP_PATTERNS:
     TOSA_QUANTIZER_SUPPORT_DICT[pattern] = ArithmeticFloatInputsCheck
+TOSA_QUANTIZER_SUPPORT_DICT[(torch.ops.aten.to.dtype,)] = CastCheck

backends/arm/test/ops/test_to_copy.py

@@ -334,33 +334,33 @@ def test_to_tosa_INT_not_delegated(test_data: Tuple):
     pipeline.run()
 
 
-_TO_COPY_QUANTIZED_IDENTITY_CAST_DATA = {
-    "int8_cast_add": lambda: (
+_TO_COPY_QUANTIZED_INT_TO_FLOAT_CAST_DATA = {
+    "int8_to_fp32_add": lambda: (
         (torch.randn(1, 3, 4, 4) * 10).to(dtype=torch.int8),
         torch.randn(1, 3, 4, 4),
         torch.float32,
     ),
-    "int16_cast_add": lambda: (
+    "int16_to_fp32_add": lambda: (
         (torch.randn(1, 3, 4, 4) * 10).to(dtype=torch.int16),
         torch.randn(1, 3, 4, 4),
         torch.float32,
     ),
-    "int32_cast_add": lambda: (
+    "int32_to_fp32_add": lambda: (
         (torch.randn(1, 3, 4, 4) * 10).to(dtype=torch.int32),
         torch.randn(1, 3, 4, 4),
         torch.float32,
     ),
 }
 
 
-_TO_COPY_QUANTIZED_IDENTITY_CAST_CAT_DATA = {
-    "int8_cast_cat": lambda: (
+_TO_COPY_QUANTIZED_INT_TO_FLOAT_CAST_CAT_DATA = {
+    "int8_to_fp32_cat": lambda: (
         (torch.randn(1, 2, 4, 4) * 10).to(dtype=torch.int8),
         torch.randn(1, 2, 4, 1),
         torch.float32,
         3,
     ),
-    "int16_cast_cat": lambda: (
+    "int16_to_fp32_cat": lambda: (
         (torch.randn(1, 2, 4, 4) * 10).to(dtype=torch.int16),
         torch.randn(1, 2, 4, 1),
         torch.float32,
@@ -369,8 +369,8 @@ def test_to_tosa_INT_not_delegated(test_data: Tuple):
 }
 
 
-@common.parametrize("test_data", _TO_COPY_QUANTIZED_IDENTITY_CAST_DATA)
-def test_to_tosa_INT_quantized_identity_cast_add(test_data: Tuple):
+@common.parametrize("test_data", _TO_COPY_QUANTIZED_INT_TO_FLOAT_CAST_DATA)
+def test_to_tosa_INT_quantized_int_to_float_cast_add(test_data: Tuple):
     x, y, new_dtype = test_data()
     pipeline = TosaPipelineINT[input_t2](
         CastAddTensor(new_dtype),
@@ -388,8 +388,8 @@ def test_to_tosa_INT_quantized_identity_cast_add(test_data: Tuple):
     pipeline.run()
 
 
-@common.parametrize("test_data", _TO_COPY_QUANTIZED_IDENTITY_CAST_CAT_DATA)
-def test_to_tosa_INT_quantized_identity_cast_cat(test_data: Tuple):
+@common.parametrize("test_data", _TO_COPY_QUANTIZED_INT_TO_FLOAT_CAST_CAT_DATA)
+def test_to_tosa_INT_quantized_int_to_float_cast_cat(test_data: Tuple):
     x, y, new_dtype, dim = test_data()
     pipeline = TosaPipelineINT[input_t2](
         CastCatTensor(new_dtype, dim),
@@ -400,8 +400,8 @@ def test_to_tosa_INT_quantized_identity_cast_cat(test_data: Tuple):
     pipeline.run()
 
 
-@common.parametrize("test_data", _TO_COPY_QUANTIZED_IDENTITY_CAST_DATA)
-def test_to_tosa_INT_quantized_identity_cast_to_unquantized_add_delegated(
+@common.parametrize("test_data", _TO_COPY_QUANTIZED_INT_TO_FLOAT_CAST_DATA)
+def test_to_tosa_INT_quantized_int_to_float_cast_to_unquantized_add_delegated(
     test_data: Tuple,
 ):
     x, y, new_dtype = test_data()
@@ -423,6 +423,46 @@ def test_to_tosa_INT_quantized_identity_cast_to_unquantized_add_delegated(
     pipeline.run()
 
 
+_TO_COPY_INT_TO_INT_CAST_DATA = {
+    "int8_to_int16": lambda: (
+        torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int8),
+        torch.int16,
+    ),
+    "int8_to_int32": lambda: (
+        torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int8),
+        torch.int32,
+    ),
+    "int16_to_int8": lambda: (
+        torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int16),
+        torch.int8,
+    ),
+    "int16_to_int32": lambda: (
+        torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int16),
+        torch.int32,
+    ),
+    "int32_to_int8": lambda: (
+        torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int32),
+        torch.int8,
+    ),
+    "int32_to_int16": lambda: (
+        torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int32),
+        torch.int16,
+    ),
+}
+
+
+@common.parametrize("test_data", _TO_COPY_INT_TO_INT_CAST_DATA)
+def test_to_tosa_INT_int_to_int_cast(test_data: Tuple):
+    test_tensor, new_dtype = test_data()
+    pipeline = TosaPipelineINT[input_t1](
+        Cast(new_dtype),
+        (test_tensor,),
+        aten_op=[],
+        exir_op=[],
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", _TO_COPY_TEST_DATA_INT)
 @common.SkipIfNoModelConverter
 def test_to_vgf_quant(test_data: Tuple):
@@ -462,6 +502,25 @@ def test_to_vgf_quant(test_data: Tuple):
     "rand_fp16_fp16": "FP16 is not supported",
 }
 
+_TO_COPY_FLOAT_IDENTITY_CAST_DATA = {
+    "fp32_to_fp32": lambda: (
+        torch.rand((1, 2, 3, 4), dtype=torch.float32),
+        torch.float32,
+    ),
+}
+
+
+@common.parametrize("test_data", _TO_COPY_FLOAT_IDENTITY_CAST_DATA)
+def test_to_tosa_INT_float_to_same_dtype_cast(test_data: Tuple):
+    test_tensor, new_dtype = test_data()
+    pipeline = TosaPipelineINT[input_t1](
+        CastAdd(new_dtype),
+        (test_tensor,),
+        aten_op=[],
+        exir_op=[],
+    )
+    pipeline.run()
+
 
 @common.parametrize(
     "test_data", _TO_COPY_TEST_DATA_REDUNDANT_CAST, xfails=redundant_xfails_FP
@@ -504,6 +563,36 @@ def test_to_tosa_INT_not_delegated_REDUNDANT_CAST(test_data: Tuple):
     pipeline.run()
 
 
+_TO_COPY_UNSUPPORTED_QUANTIZED_CAST_DATA = {
+    "fp32_to_fp16": lambda: (
+        torch.rand((1, 2, 3, 4), dtype=torch.float32),
+        torch.float16,
+    ),
+    "fp32_to_int32": lambda: (
+        torch.rand((1, 2, 3, 4), dtype=torch.float32),
+        torch.int32,
+    ),
+    "bool_to_fp32": lambda: (
+        torch.randint(0, 2, (1, 2, 3, 4), dtype=torch.bool),
+        torch.float32,
+    ),
+}
+
+
+@common.parametrize("test_data", _TO_COPY_UNSUPPORTED_QUANTIZED_CAST_DATA)
+def test_to_tosa_INT_unsupported_cast_not_delegated(test_data: Tuple):
+    test_tensor, new_dtype = test_data()
+    pipeline = OpNotSupportedPipeline[input_t1](
+        Cast(new_dtype),
+        (test_tensor,),
+        {
+            "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1
+        },
+        quantize=True,
+    )
+    pipeline.run()
+
+
 _TO_COPY_DATA_INT_U55_REJECT = {
     "rand_bool_int8": lambda: (
         torch.randint(0, 2, (1, 2, 3, 4), dtype=torch.bool),