Update on "[ET-VK] Fake u16vecn for devserver"

## Context

Copy-pasted from the newly added `maybe_fake_u16vec3` function in the codegen script:

> There is a latency benefit to using u16vecn variables to store texture position variables instead of ivecn, likely due to reduced register pressure. However, SwiftShader does not support 16 bit integer types in shaders, so this is a crude way to fallback to using ivecn to store texture positions so that testing with SwiftShader is still possible.

Differential Revision: [D65501674](https://our.internmc.facebook.com/intern/diff/D65501674/)

[ghstack-poisoned]

Author: SS-JIA

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

@@ -11,6 +11,10 @@
 import subprocess
 import sys
 import tempfile
+from datetime import datetime
+from enum import auto, Enum
+from pathlib import Path
+from typing import Any
 
 import pytest
 
@@ -19,27 +23,46 @@
 from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 
-_enabled_options: list[str] = []
+
+class arm_test_options(Enum):
+    quantize_io = auto()
+    corstone300 = auto()
+    dump_path = auto()
+    date_format = auto()
+
+
+_test_options: dict[arm_test_options, Any] = {}
 
 # ==== Pytest hooks ====
 
 
 def pytest_addoption(parser):
     parser.addoption("--arm_quantize_io", action="store_true")
     parser.addoption("--arm_run_corstone300", action="store_true")
+    parser.addoption("--default_dump_path", default=None)
+    parser.addoption("--date_format", default="%d-%b-%H:%M:%S")
 
 
 def pytest_configure(config):
     if config.option.arm_quantize_io:
         load_libquantized_ops_aot_lib()
-        _enabled_options.append("quantize_io")
+        _test_options[arm_test_options.quantize_io] = True
     if config.option.arm_run_corstone300:
         corstone300_exists = shutil.which("FVP_Corstone_SSE-300_Ethos-U55")
         if not corstone300_exists:
             raise RuntimeError(
                 "Tests are run with --arm_run_corstone300 but corstone300 FVP is not installed."
             )
-        _enabled_options.append("corstone300")
+        _test_options[arm_test_options.corstone300] = True
+    if config.option.default_dump_path:
+        dump_path = Path(config.option.default_dump_path).expanduser()
+        if dump_path.exists() and os.path.isdir(dump_path):
+            _test_options[arm_test_options.dump_path] = dump_path
+        else:
+            raise RuntimeError(
+                f"Supplied argument 'default_dump_path={dump_path}' that does not exist or is not a directory."
+            )
+    _test_options[arm_test_options.date_format] = config.option.date_format
     logging.basicConfig(level=logging.INFO, stream=sys.stdout)
 
 
@@ -54,6 +77,18 @@ def pytest_collection_modifyitems(config, items):
                 item.add_marker(skip_if_aot_lib_not_loaded)
 
 
+def pytest_sessionstart(session):
+    pass
+
+
+def pytest_sessionfinish(session, exitstatus):
+    if get_option(arm_test_options.dump_path):
+        _clean_dir(
+            get_option(arm_test_options.dump_path),
+            f"ArmTester_{get_option(arm_test_options.date_format)}.log",
+        )
+
+
 # ==== End of Pytest hooks =====
 
 
@@ -76,7 +111,9 @@ def load_libquantized_ops_aot_lib():
         torch.ops.load_library(library_path)
 
 
-def is_option_enabled(option: str, fail_if_not_enabled: bool = False) -> bool:
+def is_option_enabled(
+    option: str | arm_test_options, fail_if_not_enabled: bool = False
+) -> bool:
     """
     Returns whether an option is successfully enabled, i.e. if the flag was
     given to pytest and the necessary requirements are available.
@@ -87,7 +124,10 @@ def is_option_enabled(option: str, fail_if_not_enabled: bool = False) -> bool:
     The optional parameter 'fail_if_not_enabled' makes the function raise
       a RuntimeError instead of returning False.
     """
-    if option.lower() in _enabled_options:
+    if isinstance(option, str):
+        option = arm_test_options[option.lower()]
+
+    if option in _test_options and _test_options[option]:
         return True
     else:
         if fail_if_not_enabled:
@@ -96,6 +136,12 @@ def is_option_enabled(option: str, fail_if_not_enabled: bool = False) -> bool:
             return False
 
 
+def get_option(option: arm_test_options) -> Any | None:
+    if option in _test_options:
+        return _test_options[option]
+    return None
+
+
 def maybe_get_tosa_collate_path() -> str | None:
     """
     Checks the environment variable TOSA_TESTCASES_BASE_PATH and returns the
@@ -219,3 +265,32 @@ def get_u85_compile_spec_unbuilt(
         .dump_intermediate_artifacts_to(artifact_path)
     )
     return compile_spec
+
+
+def current_time_formated() -> str:
+    """Return current time as a formated string"""
+    return datetime.now().strftime(get_option(arm_test_options.date_format))
+
+
+def _clean_dir(dir: Path, filter: str, num_save=10):
+    sorted_files: list[tuple[datetime, Path]] = []
+    for file in dir.iterdir():
+        try:
+            creation_time = datetime.strptime(file.name, filter)
+            insert_index = -1
+            for i, to_compare in enumerate(sorted_files):
+                compare_time = to_compare[0]
+                if creation_time < compare_time:
+                    insert_index = i
+                    break
+            if insert_index == -1 and len(sorted_files) < num_save:
+                sorted_files.append((creation_time, file))
+            else:
+                sorted_files.insert(insert_index, (creation_time, file))
+        except ValueError:
+            continue
+
+    if len(sorted_files) > num_save:
+        for remove in sorted_files[0 : len(sorted_files) - num_save]:
+            file = remove[1]
+            file.unlink()

backends/arm/test/common.py

@@ -536,9 +536,6 @@ def _test_batchnorm2d_tosa_MI_pipeline(
                 compile_spec=common.get_tosa_compile_spec(),
             )
             .export()
-            .check_count(
-                {"torch.ops.aten._native_batch_norm_legit_no_training.default": 1}
-            )
             .check_not(["torch.ops.quantized_decomposed"])
             .to_edge()
             .check_count(

backends/arm/test/ops/test_batch_norm.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()