Update on "Serialize NamedData in PTE file"

1. Serialize NamedData in PTE file
2. Add NamedDataStore to EdgeProgramManager
---
Serializing NamedData is slightly different to constant/delegate data as each segment comes with its own alignment.

**An example:**
Given NamedData = {"key": data}. Data is 250 bytes.

- BackendA requires data with alignment=3
- BackendB requires data with alignment=4

Then, data0 should be serialized with alignment of lcm(3, 4) = 12

At serialization, ExecuTorch has a 'segment_alignment' that defaults to 128. Data is now serialized to lcm(12, 128) = 384.

Inside the DataSegment, we want to store the original size of the data (250). The offset of the subsequent DataSegment would be 384 bytes after the start of this one.

**Design**
Introduce a new dataclass 'AlignedData' that stores the buffer and any alignment that's required. This is used when assembling Program.segments to ensure we get lcm(buffer_alignment, segment_alignment).


Note: The default segment_alignment can be overridden inside 'ExecutorchBackendConfig'.

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

[ghstack-poisoned]

Author: lucylq

.ci/scripts/gather_test_models.py

@@ -90,7 +90,7 @@ def model_should_run_on_event(model: str, event: str) -> bool:
     We put higher priority and fast models to pull request and rest to push.
     """
     if event == "pull_request":
-        return model in ["mv3", "vit", "qwen2_5"]  # TODO: remove, just to test the ci
+        return model in ["mv3", "vit"]
     elif event == "push":
         # These are super slow. Only run it periodically
         return model not in ["dl3", "edsr", "emformer_predict"]
@@ -104,8 +104,12 @@ def model_should_run_on_target_os(model: str, target_os: str) -> bool:
     For example, a big model can be disabled in macos due to the limited macos resources.
     """
     if target_os == "macos":
+        # Disabled in macos due to limited resources, and should stay that way even if
+        # we otherwise re-enable.
         return model not in ["llava"]
-    return True
+    # Disabled globally because we have test-llava-runner-linux that does a more
+    # comprehensive E2E test of llava.
+    return model not in ["llava"]
 
 
 def export_models_for_ci() -> dict[str, dict]:

backends/arm/operator_support/__init__.py

@@ -6,7 +6,6 @@
 # pyre-unsafe
 
 from . import (  # noqa
-    bitwise_support,
     convolution_support,
     pool_2d_support,
     reduce_sum_support,

backends/arm/operator_support/bitwise_support.py

@@ -11,13 +11,13 @@
 from typing import final, Optional, Sequence, Type
 
 import torch
-
 import torch.fx as fx
+
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm._passes.fuse_quantized_activation_pass import (
     FuseQuantizedActivationPass,
 )
-from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.backends.arm.tosa_specification import Tosa_0_80, TosaSpecification
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.fx.passes.operator_support import any_chain, chain, OperatorSupportBase
 from torch.fx.passes.utils.source_matcher_utils import get_source_partitions
@@ -90,6 +90,7 @@ def tosa_support_factory(
     if not tosa_spec.support_float():
         negative_checks.append(NeedsDecompositionCheck())
         negative_checks.append(CheckProperQuantization())
+        negative_checks.append(EthosU55NotSupported(tosa_spec))
     return chain(
         any_chain(
             BaseTOSASupportList(),
@@ -111,6 +112,9 @@ def is_node_supported(
         supported = node.op == "call_function" and node.target in [
             exir_ops.edge.aten.abs.default,
             exir_ops.edge.aten.add.Tensor,
+            exir_ops.edge.aten.bitwise_and.Tensor,
+            exir_ops.edge.aten.bitwise_or.Tensor,
+            exir_ops.edge.aten.bitwise_xor.Tensor,
             exir_ops.edge.aten.expand_copy.default,
             exir_ops.edge.aten.cat.default,
             exir_ops.edge.aten.clamp.default,
@@ -170,6 +174,31 @@ def is_node_supported(
         return supported
 
 
+class EthosU55NotSupported(OperatorSupportBase):
+    """
+    Certain operators are not supported on U55. These are listed in `unsupported` in
+    is_node_supported().
+    """
+
+    def __init__(self, tosa_spec: TosaSpecification):
+        self.tosa_spec = tosa_spec
+
+    def is_node_supported(
+        self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
+    ) -> bool:
+        if isinstance(self.tosa_spec, Tosa_0_80) and self.tosa_spec.is_U55_subset:
+            unsupported_ops = [
+                exir_ops.edge.aten.bitwise_and.Tensor,
+                exir_ops.edge.aten.bitwise_or.Tensor,
+                exir_ops.edge.aten.bitwise_xor.Tensor,
+            ]
+
+            if node.target in unsupported_ops:
+                return False
+
+        return True
+
+
 class NeedsDecompositionCheck(OperatorSupportBase):
     """
     Targeted operators need to be decomposed prior to quantization in order to get a pair of q-dq-nodes surrounding

backends/arm/operator_support/tosa_supported_operators.py

@@ -1,10 +1,16 @@
 load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
+load(":targets.bzl", "define_arm_tests")
+
+
+oncall("executorch")
 
 python_library(
     name = "conftest",
     srcs = ["conftest.py"],
     deps = [
         "//executorch/exir:lib",
+        "//executorch/exir/backend:compile_spec_schema",
+        "fbsource//third-party/pypi/pytest:pytest",
     ]
 )
 
@@ -32,3 +38,33 @@ python_library(
         "fbsource//third-party/pypi/pytest:pytest",
     ]
 )
+
+python_library(
+    name = "common",
+    srcs = ["common.py"],
+    deps = [
+        ":runner_utils",
+        "//executorch/backends/xnnpack/test/tester:tester",
+        "//executorch/backends/arm:arm_backend",
+        "//executorch/exir:lib",
+        "//executorch/exir/backend:compile_spec_schema",
+        "fbsource//third-party/pypi/pytest:pytest",
+    ]
+)
+
+python_library(
+    name = "arm_tester",
+    srcs = glob(["tester/*.py"]),
+    deps = [
+        ":runner_utils",
+        ":common",
+        "//executorch/backends/arm:tosa_mapping",
+        "//executorch/backends/arm:tosa_specification",
+        "//executorch/backends/arm/quantizer:arm_quantizer",
+        "//executorch/backends/arm:arm_partitioner",
+        "//executorch/devtools/backend_debug:delegation_info",
+        "fbsource//third-party/pypi/tabulate:tabulate",
+    ]
+)
+
+define_arm_tests() 

backends/arm/test/TARGETS

@@ -13,7 +13,12 @@
 from typing import Any
 
 import pytest
-import torch
+
+try:
+    import tosa_reference_model
+except ImportError:
+    logging.warning("tosa_reference_model not found, can't run reference model tests")
+    tosa_reference_model = None
 
 """
 This file contains the pytest hooks, fixtures etc. for the Arm test suite.
@@ -24,18 +29,29 @@
 
 
 def pytest_configure(config):
-
     pytest._test_options = {}  # type: ignore[attr-defined]
-
-    if config.option.arm_run_corstoneFVP:
+    pytest._test_options["corstone_fvp"] = False  # type: ignore[attr-defined]
+    if (
+        getattr(config.option, "arm_run_corestoneFVP", False)
+        and config.option.arm_run_corstoneFVP
+    ):
         corstone300_exists = shutil.which("FVP_Corstone_SSE-300_Ethos-U55")
         corstone320_exists = shutil.which("FVP_Corstone_SSE-320")
         if not (corstone300_exists and corstone320_exists):
             raise RuntimeError(
                 "Tests are run with --arm_run_corstoneFVP but corstone FVP is not installed."
             )
+        # Only enable if we also have the TOSA reference model available.
         pytest._test_options["corstone_fvp"] = True  # type: ignore[attr-defined]
-    pytest._test_options["fast_fvp"] = config.option.fast_fvp  # type: ignore[attr-defined]
+
+    pytest._test_options["fast_fvp"] = False  # type: ignore[attr-defined]
+    if getattr(config.option, "fast_fvp", False):
+        pytest._test_options["fast_fvp"] = config.option.fast_fvp  # type: ignore[attr-defined]
+
+    # TODO: remove this flag once we have a way to run the reference model tests with Buck
+    pytest._test_options["tosa_ref_model"] = False  # type: ignore[attr-defined]
+    if tosa_reference_model is not None:
+        pytest._test_options["tosa_ref_model"] = True  # type: ignore[attr-defined]
     logging.basicConfig(level=logging.INFO, stream=sys.stdout)
 
 
@@ -44,9 +60,15 @@ def pytest_collection_modifyitems(config, items):
 
 
 def pytest_addoption(parser):
-    parser.addoption("--arm_quantize_io", action="store_true", help="Deprecated.")
-    parser.addoption("--arm_run_corstoneFVP", action="store_true")
-    parser.addoption("--fast_fvp", action="store_true")
+    def try_addoption(*args, **kwargs):
+        try:
+            parser.addoption(*args, **kwargs)
+        except Exception:
+            pass
+
+    try_addoption("--arm_quantize_io", action="store_true", help="Deprecated.")
+    try_addoption("--arm_run_corstoneFVP", action="store_true", help="Deprecated.")
+    try_addoption("--fast_fvp", action="store_true")
 
 
 def pytest_sessionstart(session):
@@ -78,6 +100,8 @@ def set_random_seed():
     Rerun with a specific seed found under a random seed test
         ARM_TEST_SEED=3478246 pytest --config-file=/dev/null --verbose -s --color=yes  backends/arm/test/ops/test_avg_pool.py -k <TESTCASE>
     """
+    import torch
+
     if os.environ.get("ARM_TEST_SEED", "RANDOM") == "RANDOM":
         random.seed()  # reset seed, in case any other test has fiddled with it
         seed = random.randint(0, 2**32 - 1)
@@ -161,6 +185,8 @@ def _load_libquantized_ops_aot_lib():
     res = subprocess.run(find_lib_cmd, capture_output=True)
     if res.returncode == 0:
         library_path = res.stdout.decode().strip()
+        import torch
+
         torch.ops.load_library(library_path)
     else:
         raise RuntimeError(

backends/arm/test/conftest.py

@@ -116,7 +116,7 @@ def _test_rescale_pipeline(
 ):
     """Tests a model with many ops that requires rescales. As more ops are quantized to int32 and
     need the InsertRescalesPass, make sure that they play nicely together."""
-    (
+    tester = (
         ArmTester(
             module,
             example_inputs=test_data,
@@ -126,8 +126,9 @@ def _test_rescale_pipeline(
         .export()
         .to_edge_transform_and_lower()
         .to_executorch()
-        .run_method_and_compare_outputs(test_data)
     )
+    if conftest.is_option_enabled("tosa_ref_model"):
+        tester.run_method_and_compare_outputs(test_data)
 
 
 def _test_rescale_pipeline_ethosu(
@@ -152,6 +153,7 @@ def _test_rescale_pipeline_ethosu(
 class TestRescales(unittest.TestCase):
 
     @parameterized.expand(RescaleNetwork.test_parameters)
+    @pytest.mark.tosa_ref_model
     def test_quantized_rescale(self, x, y):
         _test_rescale_pipeline(RescaleNetwork(), (x, y))
 

backends/arm/test/passes/test_rescale_pass.py

@@ -2,4 +2,5 @@
 addopts = --strict-markers
 markers =
     slow: Tests that take long time
-    corstone_fvp: Tests that use Corstone300 or Corstone320 FVP
+    corstone_fvp: Tests that use Corstone300 or Corstone320 FVP # And also uses TOSA reference model
+    tosa_ref_model: Tests that use TOSA reference model # Temporary!

backends/arm/test/pytest.ini

@@ -22,7 +22,6 @@
 try:
     import tosa_reference_model
 except ImportError:
-    logger.warning("tosa_reference_model not found, can't run reference model tests")
     tosa_reference_model = None
 from executorch.backends.arm.arm_backend import get_tosa_spec, is_tosa
 

backends/arm/test/runner_utils.py

@@ -0,0 +1,35 @@
+# load("//caffe2/test/fb:defs.bzl", "define_tests")
+load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest")
+load("@bazel_skylib//lib:paths.bzl", "paths")
+
+def define_arm_tests():
+    # TODO Add more tests
+    test_files = native.glob(["passes/test_*.py"])
+
+    # https://github.com/pytorch/executorch/issues/8606
+    test_files.remove("passes/test_ioquantization_pass.py")
+
+    TESTS = {}
+
+    for test_file in test_files:
+        test_file_name = paths.basename(test_file)
+        test_name = test_file_name.replace("test_", "").replace(".py", "")
+
+        python_pytest(
+            name = test_name,
+            srcs = [test_file],
+            pytest_config = "pytest.ini",
+            resources = ["conftest.py"],
+            compile = "with-source",
+            typing = False,
+            preload_deps = [
+                "//executorch/kernels/quantized:custom_ops_generated_lib",
+            ],
+            deps = [
+                ":arm_tester",
+                ":conftest",
+                "//executorch/exir:lib",
+                "fbsource//third-party/pypi/pytest:pytest",
+                "fbsource//third-party/pypi/parameterized:parameterized",
+            ],
+        )