Add cat/stack ops to generic annotator

Differential Revision: D65067790

Pull Request resolved: pytorch#6494

Author: SaoirseARM

backends/arm/quantizer/arm_quantizer.py

@@ -268,7 +268,6 @@ class ArmQuantizer(Quantizer):
         "sub",
         "mul",
         "mm",
-        "cat",
         "one_to_one",
         "generic",
         "sum",

backends/arm/quantizer/arm_quantizer_utils.py

@@ -144,21 +144,10 @@ def is_share_obs_or_fq_op(op: Callable) -> bool:
         torch.ops.aten.mean.dim,
         torch.ops.aten.permute.default,
         torch.ops.aten.permute_copy.default,
-        torch.ops.aten.squeeze.dim,
-        torch.ops.aten.squeeze.dims,
-        torch.ops.aten.squeeze.default,
-        torch.ops.aten.squeeze_copy.dim,
-        torch.ops.aten.unsqueeze.default,
-        torch.ops.aten.unsqueeze_copy.default,
         # TODO: remove?
         torch.ops.aten.adaptive_avg_pool2d.default,
         torch.ops.aten.avg_pool2d.default,
-        torch.ops.aten.view_copy.default,
-        torch.ops.aten.view.default,
         torch.ops.aten.full.default,
-        torch.ops.aten.slice.Tensor,
-        torch.ops.aten.split.Tensor,
-        torch.ops.aten.split_with_sizes.default,
         torch.ops.aten.flatten.using_ints,
         torch.ops.aten.dropout.default,
         operator.getitem,

backends/arm/quantizer/quantization_annotation/__init__.py

@@ -51,7 +51,6 @@ def decorator(annotator: AnnotatorType):
 from . import (  # noqa
     adaptive_ang_pool2d_annotator,
     add_annotator,
-    cat_annotator,
     conv_annotator,
     generic_annotator,
     linear_annotator,

backends/arm/quantizer/quantization_annotation/cat_annotator.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-unsafe
-
 from typing import Callable, List, Optional
 
 import torch
@@ -24,6 +23,9 @@
     # DATA LAYOUT OPS
     torch.ops.aten.squeeze.default,
     torch.ops.aten.squeeze_copy.default,
+    torch.ops.aten.squeeze_copy.dim,
+    torch.ops.aten.squeeze.dim,
+    torch.ops.aten.squeeze.dims,
     torch.ops.aten.unsqueeze.default,
     torch.ops.aten.unsqueeze_copy.default,
     torch.ops.aten.reshape.default,
@@ -33,19 +35,21 @@
     # torch.ops.aten.view_as_complex_copy.default,
     # torch.ops.aten.view_as_real.default,
     # torch.ops.aten.view_as_real_copy.default,
+    torch.ops.aten.view.default,
     torch.ops.aten.view_copy.default,
     torch.ops.aten.select.int,
     torch.ops.aten.select_copy.int,
     torch.ops.aten.slice.Tensor,
     torch.ops.aten.slice_copy.Tensor,
-    # 'concat' should be handled separately as it has a sequence of inputs and
-    # makes the implementation unnecessary complicated.
-    # torch.ops.aten.concat.default,
+    torch.ops.aten.split.Tensor,
+    torch.ops.aten.split_with_sizes.default,
     torch.ops.aten.transpose.Dimname,
     torch.ops.aten.transpose.int,
     torch.ops.aten.transpose_copy.int,
     torch.ops.aten.tile.default,
     torch.ops.aten.flip.default,
+    torch.ops.aten.cat.default,
+    torch.ops.aten.stack.default,
 ]
 
 
@@ -66,15 +70,31 @@ def _annotate_generic(
         if arm_quantizer_utils.is_annotated(node):
             continue
 
-        input_node = node.args[0]
+        input_acts = node.args[0]
+
+        # Check to see if there are multiple inputs.
+        # this allows for stack/cat ops to be annotated
+        # in a similar way.
+        has_multi_inputs = isinstance(input_acts, list)
+
+        input_act0 = input_acts[0] if has_multi_inputs else input_acts
 
         # Using a non-shared quantization spec here as a SharedQuantizationSpec
         # can lead to a recursion.
         _annotate_input_qspec_map(
-            node, input_node, quantization_config.get_input_act_qspec()
+            node, input_act0, quantization_config.get_input_act_qspec()
         )
-        _annotate_output_qspec(node, SharedQuantizationSpec((input_node, node)))
+        shared_with_input0_qspec = SharedQuantizationSpec((input_act0, node))
+
+        if has_multi_inputs:
+            # For the rest of the inputs, share qspec with first.
+            for input_act in input_acts[1:]:
+                if input_act is not input_act0:
+                    node.meta["quantization_annotation"].input_qspec_map[
+                        input_act
+                    ] = shared_with_input0_qspec
 
+        _annotate_output_qspec(node, shared_with_input0_qspec)
         arm_quantizer_utils.mark_nodes_as_annotated([node])
         annotated_partitions.append([node])
 

backends/arm/quantizer/quantization_annotation/generic_annotator.py

@@ -8,13 +8,9 @@
 from typing import Tuple
 
 import torch
-from executorch.backends.arm.quantizer.arm_quantizer import (
-    ArmQuantizer,
-    get_symmetric_quantization_config,
-)
+
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.arm_tester import ArmTester
-from executorch.backends.xnnpack.test.tester.tester import Quantize
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from parameterized import parameterized
 
@@ -59,7 +55,6 @@ def _test_slice_tosa_BI_pipeline(
         self, module: torch.nn.Module, test_data: Tuple[torch.Tensor], permute: bool
     ):
 
-        quantizer = ArmQuantizer().set_io(get_symmetric_quantization_config())
         (
             ArmTester(
                 module,
@@ -68,7 +63,7 @@ def _test_slice_tosa_BI_pipeline(
                     permute_memory_to_nhwc=permute
                 ),
             )
-            .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
+            .quantize()
             .export()
             .check(["torch.ops.aten.slice.Tensor"])
             .to_edge()
@@ -84,14 +79,13 @@ def _test_slice_ethos_BI_pipeline(
         module: torch.nn.Module,
         test_data: Tuple[torch.Tensor],
     ):
-        quantizer = ArmQuantizer().set_io(get_symmetric_quantization_config())
         (
             ArmTester(
                 module,
                 example_inputs=test_data,
                 compile_spec=common.get_u55_compile_spec(),
             )
-            .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
+            .quantize()
             .export()
             .check(["torch.ops.aten.slice.Tensor"])
             .to_edge()

backends/arm/test/ops/test_slice.py

@@ -7,13 +7,9 @@
 import unittest
 
 import torch
-from executorch.backends.arm.quantizer.arm_quantizer import (
-    ArmQuantizer,
-    get_symmetric_quantization_config,
-)
+
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.arm_tester import ArmTester
-from executorch.backends.xnnpack.test.tester.tester import Quantize
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from parameterized import parameterized
 
@@ -79,14 +75,13 @@ def _test_split_tosa_BI_pipeline(
         self, module: torch.nn.Module, test_data: test_data_t
     ):
 
-        quantizer = ArmQuantizer().set_io(get_symmetric_quantization_config())
         (
             ArmTester(
                 module,
                 example_inputs=test_data,
                 compile_spec=common.get_tosa_compile_spec(),
             )
-            .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
+            .quantize()
             .export()
             .to_edge()
             .partition()
@@ -98,14 +93,13 @@ def _test_split_tosa_BI_pipeline(
     def _test_split_ethosu_BI_pipeline(
         self, compile_spec: CompileSpec, module: torch.nn.Module, test_data: test_data_t
     ):
-        quantizer = ArmQuantizer().set_io(get_symmetric_quantization_config())
         (
             ArmTester(
                 module,
                 example_inputs=test_data,
                 compile_spec=compile_spec,
             )
-            .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
+            .quantize()
             .export()
             .check(["torch.ops.aten.split.Tensor"])
             .to_edge()

backends/arm/test/ops/test_split.py

@@ -13,14 +13,9 @@
 
 import torch
 
-from executorch.backends.arm.quantizer.arm_quantizer import (
-    ArmQuantizer,
-    get_symmetric_quantization_config,
-)
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.arm_tester import ArmTester
 
-from executorch.backends.xnnpack.test.tester.tester import Quantize
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from parameterized import parameterized
 
@@ -83,14 +78,13 @@ def _test_squeeze_tosa_BI_pipeline(
         test_data: Tuple[torch.Tensor, Optional[tuple[int]]],
         export_target: str,
     ):
-        quantizer = ArmQuantizer().set_io(get_symmetric_quantization_config())
         (
             ArmTester(
                 module,
                 example_inputs=test_data,
                 compile_spec=common.get_tosa_compile_spec(),
             )
-            .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
+            .quantize()
             .export()
             .check_count({export_target: 1})
             .to_edge()
@@ -107,10 +101,9 @@ def _test_squeeze_ethosu_BI_pipeline(
         test_data: Tuple[torch.Tensor, Optional[tuple[int]]],
         export_target: str,
     ):
-        quantizer = ArmQuantizer().set_io(get_symmetric_quantization_config())
         (
             ArmTester(module, example_inputs=test_data, compile_spec=compile_spec)
-            .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
+            .quantize()
             .export()
             .check_count({export_target: 1})
             .to_edge()

backends/arm/test/ops/test_squeeze.py

@@ -13,14 +13,9 @@
 
 import torch
 
-from executorch.backends.arm.quantizer.arm_quantizer import (
-    ArmQuantizer,
-    get_symmetric_quantization_config,
-)
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.arm_tester import ArmTester
 
-from executorch.backends.xnnpack.test.tester.tester import Quantize
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from parameterized import parameterized
 
@@ -54,14 +49,13 @@ def _test_unsqueeze_tosa_MI_pipeline(
     def _test_unsqueeze_tosa_BI_pipeline(
         self, module: torch.nn.Module, test_data: Tuple[torch.Tensor, int]
     ):
-        quantizer = ArmQuantizer().set_io(get_symmetric_quantization_config())
         (
             ArmTester(
                 module,
                 example_inputs=test_data,
                 compile_spec=common.get_tosa_compile_spec(),
             )
-            .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
+            .quantize()
             .export()
             .check_count({"torch.ops.aten.unsqueeze.default": 1})
             .to_edge()
@@ -77,14 +71,13 @@ def _test_unsqueeze_ethosu_BI_pipeline(
         module: torch.nn.Module,
         test_data: Tuple[torch.Tensor, int],
     ):
-        quantizer = ArmQuantizer().set_io(get_symmetric_quantization_config())
         (
             ArmTester(
                 module,
                 example_inputs=test_data,
                 compile_spec=compile_spec,
             )
-            .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
+            .quantize()
             .export()
             .check_count({"torch.ops.aten.unsqueeze.default": 1})
             .to_edge()