Update on "move rope related logic together"

Right now, rope related code scatters around a few different places in `llama_transformer`. It makes it hard to make changes to rope related things.

This PR moves all rope related logic into its own module.

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

[ghstack-poisoned]

Author: helunwencser

backends/arm/TARGETS

@@ -1,3 +1,4 @@
+# @noautodeps
 load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
 
 python_library(

backends/arm/operators/TARGETS

@@ -1,3 +1,4 @@
+# @noautodeps
 load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
 
 python_library(

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(