Arm backend: Support quantizing no-op fp casts.

If a int to fp cast goes directly into a quantized
operator, the cast is a noop if it is quantized
with unit quantization parameters (scale=1, zp=0).
This means it can be handled, even if the backend
doesn't support fp.

- Add Fixed unit qparam annotation to such a cast.
- Allow it in the partitioner.
- Finally, make sure the cast dtype is correct after
folding quant nodes.

This pattern can appear in the wild in a model,
or in a decomposition where you want to quantize
an integer tensor.

Signed-off-by: Erik Lundell <erik.lundell@arm.com>
Change-Id: I93bfd17ba51e25f121f61cbe3003e9c6b9891401

Author: Erik-Lundell

backends/arm/_passes/decompose_div_tensor_mode.py

@@ -46,7 +46,7 @@
     "logical_xor": torch.ops.aten.logical_xor.default,
     "mul": torch.ops.aten.mul.Tensor,
     "sub": torch.ops.aten.sub.Tensor,
-    "to": torch.ops.dim_order_ops._to_dim_order_copy.default,
+    "to": torch.ops.aten.to.dtype,
     "where": torch.ops.aten.where.self,
 }
 

backends/arm/_passes/fold_qdq_with_annotated_qparams_pass.py

@@ -305,6 +305,25 @@ def is_foldable(node: Node) -> bool:
             return False
         return True
 
+    @staticmethod
+    def _correct_output_dtype(node: torch.fx.Node):
+        if node.target not in {
+            exir_ops.edge.aten.sum.dim_IntList,
+            exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+        }:
+            return
+        if len(node.meta["output_qparams"]) == 0:
+            return
+        output_qparams = cast(QuantArgs, node.meta["output_qparams"][0])
+
+        if node.target == exir_ops.edge.dim_order_ops._to_dim_order_copy.default:
+            if output_qparams.scale != 1.0 or output_qparams.zp != 0.0:
+                raise ValueError(
+                    f"Expected quantized {node.target} '{node.name}' to have unit scale and zero point."
+                )
+
+        set_node_arg(node, "dtype", output_qparams.dtype)
+
     def call(self, graph_module: GraphModule) -> PassResult:  # noqa: C901
 
         # Loop over the graph nodes and find any node in the 'targeted_ops' list.
@@ -355,13 +374,7 @@ def call(self, graph_module: GraphModule) -> PassResult:  # noqa: C901
 
             # Some op(s) contain a "dtype" key in their node kwargs. Set this
             # to the type of output qparams.
-            output_qparams = n.meta["output_qparams"]
-            if (
-                n.target in {exir_ops.edge.aten.sum.dim_IntList}
-                and len(output_qparams) > 0
-            ):
-                output_dtype = output_qparams[0].dtype
-                set_node_arg(n, "dtype", output_dtype)
+            FoldAndAnnotateQParamsPass._correct_output_dtype(n)
 
             if n.target in (
                 torch.ops.higher_order.cond,

backends/arm/operator_support/to_dim_order_copy_support.py

@@ -139,6 +139,20 @@ def _merge_supported_types(
         torch.float8_e5m2: [torch.bfloat16],
     }
 
+    @staticmethod
+    def _is_quantized_identity_cast(node: torch.fx.Node) -> bool:
+        for user in node.users:
+            if (
+                not user.target
+                == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
+            ):
+                return False
+            scale = user.args[1]
+            zp = user.args[2]
+            if scale != 1.0 or zp != 0.0:
+                return False
+        return True
+
     def is_node_tosa_supported(  # noqa: C901
         self, node: fx.Node, tosa_spec: TosaSpecification
     ) -> bool:
@@ -228,6 +242,11 @@ def is_node_tosa_supported(  # noqa: C901
             )
             return False
         if output_val.dtype not in supported_dtypes[input_dtype]:
+            if (
+                tosa_spec.support_integer()
+                and ToCopySupported._is_quantized_identity_cast(node)
+            ):
+                return True
             self.reporter.report_reject(
                 node,
                 (

backends/arm/quantizer/quantization_annotator.py

@@ -956,6 +956,27 @@ def any_or_hardtanh_min_zero(n: Node):
         shared_qspec = SharedQuantizationSpec(input_node)
         quant_properties.quant_inputs = [_QuantProperty(0, shared_qspec)]
         quant_properties.quant_output = _QuantProperty(0, shared_qspec)
+    elif node.target == torch.ops.aten.to.dtype:
+        # If we quantize a cast(fp32) with unit scale and same dtype as input, we can handle it as a no-op in the backend.
+        input_val = node.all_input_nodes[0].meta.get("val", None)
+        if input_val is None:
+            return None
+
+        if input_val.dtype not in (torch.int8, torch.int16, torch.int32):
+            return None
+
+        quant_properties.quant_output = _QuantProperty(
+            0,
+            FixedQParamsQuantizationSpec(
+                dtype=input_val.dtype,
+                scale=1.0,
+                zero_point=0,
+                quant_max=torch.iinfo(input_val.dtype).max,
+                quant_min=torch.iinfo(input_val.dtype).min,
+                qscheme=torch.per_tensor_symmetric,
+                is_dynamic=False,
+            ),
+        )
     elif node.target in (
         torch.ops.higher_order.cond,
         torch.ops.higher_order.while_loop,

backends/arm/test/ops/test_div_tensor_mode.py

@@ -172,7 +172,6 @@ def test_div_tensor_mode_u85_INT(data):
         exir_ops=[],
         use_to_edge_transform_and_lower=True,
     )
-    pipeline.tester.use_portable_ops = True
     pipeline.run()
 
 

backends/arm/test/ops/test_to_copy.py

@@ -20,6 +20,7 @@
 )
 
 input_t1 = Tuple[torch.Tensor]  # Input x
+input_t2 = Tuple[torch.Tensor, torch.Tensor]  # Input x, y
 
 
 class Cast(torch.nn.Module):
@@ -40,6 +41,40 @@ def forward(self, x: torch.Tensor):
         return x.to(dtype=self.target_dtype) + x.to(dtype=self.target_dtype)
 
 
+class CastAddTensor(torch.nn.Module):
+    def __init__(self, target_dtype):
+        super().__init__()
+        self.target_dtype = target_dtype
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        return x.to(dtype=self.target_dtype) + y
+
+
+class AddModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        return x + y
+
+
+class CastToAddModule(torch.nn.Module):
+    def __init__(self, target_dtype):
+        super().__init__()
+        self.target_dtype = target_dtype
+        self.add = AddModule()
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        return self.add(x.to(dtype=self.target_dtype), y)
+
+
+class CastCatTensor(torch.nn.Module):
+    def __init__(self, target_dtype, dim: int):
+        super().__init__()
+        self.target_dtype = target_dtype
+        self.dim = dim
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        return torch.cat((x.to(dtype=self.target_dtype), y), dim=self.dim)
+
+
 """
 Tests the _to_copy operation.
 
@@ -262,14 +297,6 @@ def test_to_vgf_no_quant(test_data: Tuple):
 in ToCopySupported::is_node_tosa_supported() before it goes into the delegated graph.
 """
 _TO_COPY_TEST_DATA_INT = {
-    "rand_int8_fp32": lambda: (
-        torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int8),
-        torch.float32,
-    ),
-    "rand_int16_fp32": lambda: (
-        torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int16),
-        torch.float32,
-    ),
     "rand_int32_fp32": lambda: (
         torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int32),
         torch.float32,
@@ -300,6 +327,95 @@ def test_to_tosa_INT_not_delegated(test_data: Tuple):
     pipeline.run()
 
 
+_TO_COPY_QUANTIZED_IDENTITY_CAST_DATA = {
+    "int8_cast_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: (
+        (torch.randn(1, 3, 4, 4) * 10).to(dtype=torch.int16),
+        torch.randn(1, 3, 4, 4),
+        torch.float32,
+    ),
+    "int32_cast_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: (
+        (torch.randn(1, 2, 4, 4) * 10).to(dtype=torch.int8),
+        torch.randn(1, 2, 4, 1),
+        torch.float32,
+        3,
+    ),
+    "int16_cast_cat": lambda: (
+        (torch.randn(1, 2, 4, 4) * 10).to(dtype=torch.int16),
+        torch.randn(1, 2, 4, 1),
+        torch.float32,
+        3,
+    ),
+}
+
+
+@common.parametrize("test_data", _TO_COPY_QUANTIZED_IDENTITY_CAST_DATA)
+def test_to_tosa_INT_quantized_identity_cast_add(test_data: Tuple):
+    x, y, new_dtype = test_data()
+    pipeline = TosaPipelineINT[input_t2](
+        CastAddTensor(new_dtype),
+        (x, y),
+        aten_op=["torch.ops.aten.add.Tensor"],
+        exir_op=["executorch_exir_dialects_edge__ops_aten_add_Tensor"],
+        qtol=1,
+    )
+    pipeline.change_args(
+        "check_count.exir",
+        {
+            "torch.ops.higher_order.executorch_call_delegate": 1,
+        },
+    )
+    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):
+    x, y, new_dtype, dim = test_data()
+    pipeline = TosaPipelineINT[input_t2](
+        CastCatTensor(new_dtype, dim),
+        (x, y),
+        aten_op=["torch.ops.aten.cat.default"],
+        exir_op=["executorch_exir_dialects_edge__ops_aten_cat_default"],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", _TO_COPY_QUANTIZED_IDENTITY_CAST_DATA)
+def test_to_tosa_INT_quantized_identity_cast_to_unquantized_add_delegated(
+    test_data: Tuple,
+):
+    x, y, new_dtype = test_data()
+    pipeline = TosaPipelineINT[input_t2](
+        CastToAddModule(new_dtype),
+        (x, y),
+        aten_op=["torch.ops.aten.add.Tensor"],
+        exir_op=["executorch_exir_dialects_edge__ops_aten_add_Tensor"],
+    )
+    pipeline.quantizer.set_module_name("add", None)
+    pipeline.pop_stage("check_not.exir")
+    pipeline.change_args(
+        "check_count.exir",
+        {
+            "torch.ops.higher_order.executorch_call_delegate": 1,
+            "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 0,
+        },
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", _TO_COPY_TEST_DATA_INT)
 @common.SkipIfNoModelConverter
 def test_to_vgf_quant(test_data: Tuple):