Move compile spec to ArmTester interface

* Create compile spec builder
* Added default compile spec for unit tests

Change-Id: Ieae15b213969c703fc0bfcfd99a6a84ae3412676
Signed-off-by: Fredrik Knutsson <[email protected]>
Co-authored-by: Per Åstrand <[email protected]>

Author: freddan80

backends/arm/arm_backend.py

@@ -38,6 +38,118 @@
     logger.setLevel(logging.INFO)
 
 
+class ArmCompileSpecBuilder:
+    def __init__(self):
+        self.compile_spec: List[CompileSpec] = []
+        self.compiler_flags = []
+        self.output_format = None
+        self.path_for_intermediates = None
+        self.permute_nhwc = False
+
+    def ethosu_compile_spec(
+        self,
+        config: str,
+        system_config: Optional[str] = None,
+        memory_mode: Optional[str] = None,
+        extra_flags: Optional[str] = None,
+        config_ini: Optional[str] = "Arm/vela.ini",
+    ):
+        """
+        Generate compile spec for Ethos-U NPU
+
+        Args:
+            config: Ethos-U accelerator configuration, e.g. ethos-u55-128
+            system_config: System configuration to select from the Vel
+                configuration file
+            memory_mode: Memory mode to select from the Vela configuration file
+            extra_flags: Extra flags for the Vela compiler
+            config_ini: Vela configuration file(s) in Python ConfigParser .ini
+                file format
+        """
+        assert (
+            self.output_format is None
+        ), f"Output format already set to f{self.output_format}"
+        self.output_format = "vela"
+        self.compiler_flags = [
+            f"--accelerator-config={config}",
+            f"--config={config_ini}",
+        ]
+        if system_config is not None:
+            self.compiler_flags.append(f"--system-config={system_config}")
+        if memory_mode is not None:
+            self.compiler_flags.append(f"--memory-mode={memory_mode}")
+        if extra_flags is not None:
+            self.compiler_flags.append(extra_flags)
+
+        return self
+
+    def tosa_compile_spec(self):
+        """
+        Generate compile spec for TOSA flatbuffer output
+        """
+        assert (
+            self.output_format is None
+        ), f"Output format already set: {self.output_format}"
+        self.output_format = "tosa"
+        return self
+
+    def dump_intermediate_tosa(self, output_path: str):
+        """
+        Output intermediate .tosa file
+        """
+        self.path_for_intermediates = output_path
+        return self
+
+    def set_permute_memory_format(self, set_nhwc_permutation: bool = True):
+        self.permute_nhwc = set_nhwc_permutation
+        return self
+
+    def build(self):
+        """
+        Generate a list of compile spec objects from the builder
+        """
+        if self.output_format == "vela":
+            self.compile_spec += [
+                CompileSpec("output_format", "vela".encode()),
+                CompileSpec("compile_flags", " ".join(self.compiler_flags).encode()),
+            ]
+        elif self.output_format == "tosa":
+            self.compile_spec.append(CompileSpec("output_format", "tosa".encode()))
+
+        if self.path_for_intermediates is not None:
+            self.compile_spec.append(
+                CompileSpec("debug_tosa_path", self.path_for_intermediates.encode())
+            )
+
+        if self.permute_nhwc:
+            self.compile_spec.append(
+                CompileSpec("permute_memory_format", "nhwc".encode())
+            )
+
+        return self.compile_spec
+
+
+def is_permute_memory(compile_spec: List[CompileSpec]) -> bool:
+    for spec in compile_spec:
+        if spec.key == "permute_memory_format":
+            return spec.value.decode() == "nhwc"
+    return False
+
+
+def is_tosa(compile_spec: List[CompileSpec]) -> bool:
+    for spec in compile_spec:
+        if spec.key == "output_format":
+            return spec.value.decode() == "tosa"
+    return False
+
+
+def get_intermediate_path(compile_spec: List[CompileSpec]) -> str:
+    for spec in compile_spec:
+        if spec.key == "debug_tosa_path":
+            return spec.value.decode()
+    return None
+
+
 def generate_ethosu_compile_spec(
     config: str,
     permute_memory_to_nhwc: Optional[bool] = None,
@@ -46,45 +158,31 @@ def generate_ethosu_compile_spec(
     extra_flags: Optional[str] = None,
     config_ini: Optional[str] = "Arm/vela.ini",
 ) -> List[CompileSpec]:
-    """
-    Generate compile spec for Ethos-U NPU
-    """
-    compiler_flags = [f"--accelerator-config={config}", f"--config={config_ini}"]
-    if system_config is not None:
-        compiler_flags.append(f"--system-config={system_config}")
-    if memory_mode is not None:
-        compiler_flags.append(f"--memory-mode={memory_mode}")
-    if extra_flags is not None:
-        compiler_flags.append(extra_flags)
-
-    compile_spec = [
-        CompileSpec("output_format", "vela".encode()),
-        CompileSpec("compile_flags", " ".join(compiler_flags).encode()),
-    ]
-
-    if permute_memory_to_nhwc:
-        compile_spec.append(CompileSpec("permute_memory_format", "nhwc".encode()))
-
-    return compile_spec
+    return (
+        ArmCompileSpecBuilder()
+        .ethosu_compile_spec(
+            config,
+            system_config=system_config,
+            memory_mode=memory_mode,
+            extra_flags=extra_flags,
+            config_ini=config_ini,
+        )
+        .set_permute_memory_format(permute_memory_to_nhwc)
+        .build()
+    )
 
 
 def generate_tosa_compile_spec(
     permute_memory_to_nhwc: Optional[bool] = None,
     output_path: Optional[str] = None,
 ) -> List[CompileSpec]:
-    """
-    Generate compile spec for TOSA flatbuffer output
-    """
-
-    compile_spec = [CompileSpec("output_format", "tosa".encode())]
-
-    if permute_memory_to_nhwc:
-        compile_spec.append(CompileSpec("permute_memory_format", "nhwc".encode()))
-
-    if output_path is not None:
-        compile_spec.append(CompileSpec("debug_tosa_path", output_path.encode()))
-
-    return compile_spec
+    return (
+        ArmCompileSpecBuilder()
+        .tosa_compile_spec()
+        .set_permute_memory_format(permute_memory_to_nhwc)
+        .dump_intermediate_tosa(output_path)
+        .build()
+    )
 
 
 @final

backends/arm/test/common.py

@@ -5,6 +5,9 @@
 # LICENSE file in the root directory of this source tree.
 
 import shutil
+import tempfile
+
+from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder
 
 # TODO: fixme! These globs are a temporary workaround. Reasoning:
 # Running the jobs in _unittest.yml will not work since that environment doesn't
@@ -13,3 +16,36 @@
 # should be installed in the CI env.
 TOSA_REF_MODEL_INSTALLED = shutil.which("tosa_reference_model")
 VELA_INSTALLED = shutil.which("vela")
+
+
+def get_tosa_compile_spec(permute_memory_to_nhwc=False, custom_path=None):
+    """
+    Default compile spec for TOSA tests.
+    """
+    intermediate_path = custom_path or tempfile.mkdtemp(prefix="arm_tosa_")
+    compile_spec = (
+        ArmCompileSpecBuilder()
+        .tosa_compile_spec()
+        .set_permute_memory_format(permute_memory_to_nhwc)
+        .dump_intermediate_tosa(intermediate_path)
+        .build()
+    )
+    return compile_spec
+
+
+def get_u55_compile_spec(permute_memory_to_nhwc=False):
+    """
+    Default compile spec for Ethos-U55 tests.
+    """
+    compile_spec = (
+        ArmCompileSpecBuilder()
+        .ethosu_compile_spec(
+            "ethos-u55-128",
+            system_config="Ethos_U55_High_End_Embedded",
+            memory_mode="Shared_Sram",
+            extra_flags=None,
+        )
+        .set_permute_memory_format(permute_memory_to_nhwc)
+        .build()
+    )
+    return compile_spec

backends/arm/test/misc/test_debug_feats.py

@@ -10,8 +10,9 @@
 import unittest
 
 import torch
-from executorch.backends.arm.test.test_models import TosaProfile
-from executorch.backends.arm.test.tester.arm_tester import ArmBackendSelector, ArmTester
+from executorch.backends.arm.test import common
+
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -45,8 +46,7 @@ def _tosa_MI_pipeline(self, module: torch.nn.Module, dump_file=None):
             ArmTester(
                 module,
                 inputs=module.get_inputs(),
-                profile=TosaProfile.MI,
-                backend=ArmBackendSelector.TOSA,
+                compile_spec=common.get_tosa_compile_spec(),
             )
             .export()
             .to_edge()
@@ -60,8 +60,7 @@ def _tosa_BI_pipeline(self, module: torch.nn.Module, dump_file=None):
             ArmTester(
                 module,
                 inputs=module.get_inputs(),
-                profile=TosaProfile.BI,
-                backend=ArmBackendSelector.TOSA,
+                compile_spec=common.get_tosa_compile_spec(),
             )
             .quantize()
             .export()

backends/arm/test/models/test_mobilenet_v2_arm.py

@@ -11,8 +11,8 @@
 import torch
 import torchvision.models as models
 from executorch.backends.arm.test import common
-from executorch.backends.arm.test.test_models import TosaProfile
-from executorch.backends.arm.test.tester.arm_tester import ArmBackendSelector, ArmTester
+
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
 from executorch.backends.xnnpack.test.tester.tester import Quantize
 from torchvision.models.mobilenetv2 import MobileNet_V2_Weights
 
@@ -46,9 +46,7 @@ def test_mv2_tosa_MI(self):
             ArmTester(
                 self.mv2,
                 inputs=self.model_inputs,
-                profile=TosaProfile.MI,
-                backend=ArmBackendSelector.TOSA,
-                permute_memory_to_nhwc=True,
+                compile_spec=common.get_tosa_compile_spec(permute_memory_to_nhwc=True),
             )
             .export()
             .to_edge()
@@ -62,9 +60,7 @@ def test_mv2_tosa_BI(self):
             ArmTester(
                 self.mv2,
                 inputs=self.model_inputs,
-                profile=TosaProfile.BI,
-                backend=ArmBackendSelector.TOSA,
-                permute_memory_to_nhwc=True,
+                compile_spec=common.get_tosa_compile_spec(permute_memory_to_nhwc=True),
             )
             .quantize(Quantize(calibrate=False))
             .export()
@@ -74,7 +70,7 @@ def test_mv2_tosa_BI(self):
             .to_executorch()
         )
         if common.TOSA_REF_MODEL_INSTALLED:
-            tester.run_method().compare_outputs()
+            tester.run_method_and_compare_outputs()
         else:
             logger.warning(
                 "TOSA ref model tool not installed, skip numerical correctness tests"
@@ -89,9 +85,7 @@ def test_mv2_u55_BI(self):
             ArmTester(
                 self.mv2,
                 inputs=self.model_inputs,
-                profile=TosaProfile.BI,
-                backend=ArmBackendSelector.ETHOS_U55,
-                permute_memory_to_nhwc=True,
+                compile_spec=common.get_u55_compile_spec(permute_memory_to_nhwc=True),
             )
             .quantize(Quantize(calibrate=False))
             .export()

backends/arm/test/ops/test_add.py

@@ -12,8 +12,8 @@
 
 import torch
 from executorch.backends.arm.test import common
-from executorch.backends.arm.test.test_models import TosaProfile
-from executorch.backends.arm.test.tester.arm_tester import ArmBackendSelector, ArmTester
+
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
 from parameterized import parameterized
 
 logger = logging.getLogger(__name__)
@@ -57,9 +57,7 @@ def _test_add_tosa_MI_pipeline(
             ArmTester(
                 module,
                 inputs=test_data,
-                profile=TosaProfile.MI,
-                backend=ArmBackendSelector.TOSA,
-                permute_memory_to_nhwc=False,
+                compile_spec=common.get_tosa_compile_spec(),
             )
             .export()
             .check_count({"torch.ops.aten.add.Tensor": 1})
@@ -70,7 +68,7 @@ def _test_add_tosa_MI_pipeline(
             .to_executorch()
         )
         if common.TOSA_REF_MODEL_INSTALLED:
-            tester.run_method().compare_outputs()
+            tester.run_method_and_compare_outputs()
         else:
             logger.warning(
                 "TOSA ref model tool not installed, skip numerical correctness tests"
@@ -83,9 +81,7 @@ def _test_add_tosa_BI_pipeline(
             ArmTester(
                 module,
                 inputs=test_data,
-                profile=TosaProfile.BI,
-                backend=ArmBackendSelector.TOSA,
-                permute_memory_to_nhwc=False,
+                compile_spec=common.get_tosa_compile_spec(),
             )
             .quantize()
             .export()
@@ -98,7 +94,7 @@ def _test_add_tosa_BI_pipeline(
         )
 
         if common.TOSA_REF_MODEL_INSTALLED:
-            tester.run_method().compare_outputs(qtol=1)
+            tester.run_method_and_compare_outputs(qtol=1)
         else:
             logger.warning(
                 "TOSA ref model tool not installed, skip numerical correctness tests"
@@ -111,8 +107,7 @@ def _test_add_u55_BI_pipeline(
             ArmTester(
                 module,
                 inputs=test_data,
-                profile=TosaProfile.BI,
-                backend=ArmBackendSelector.ETHOS_U55,
+                compile_spec=common.get_u55_compile_spec(),
             )
             .quantize()
             .export()