Support to add running stages for E2E test

Co-authored-by: chong-chen <[email protected]>

Author: Jiseong-oh

backends/samsung/test/tester/samsung_tester.py

@@ -4,51 +4,80 @@
 # Licensed under the BSD License (the "License"); you may not use this file
 # except in compliance with the License. See the license file in the root
 # directory of this source tree for more details.
-from typing import List, Optional, Tuple
+import copy
+from typing import Any, List, Optional, Sequence, Tuple, Union
 
 import executorch.backends.test.harness.stages as BaseStages
 import torch
 from executorch.backends.samsung.partition.enn_partitioner import EnnPartitioner
-from executorch.backends.samsung.utils.export_utils import get_edge_compile_config
-
+from executorch.backends.samsung.quantizer.quantizer import EnnQuantizer, Precision
+from executorch.backends.samsung.test.utils import RuntimeExecutor
+from executorch.backends.samsung.utils.export_utils import (
+    get_edge_compile_config,
+    get_enn_pass_list,
+)
 from executorch.backends.test.harness import Tester as TesterBase
+from executorch.backends.test.harness.stages import StageType
 from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
 from executorch.exir.backend.backend_details import CompileSpec
 
+from executorch.exir.pass_manager import PassType
 from torch.export import ExportedProgram
 
+from torchao.quantization.pt2e.quantizer import Quantizer
+
 
 class Export(BaseStages.Export):
     pass
 
 
 class Quantize(BaseStages.Quantize):
-    pass
+    def __init__(
+        self,
+        quantizer: Optional[Quantizer] = None,
+        quantization_config: Optional[Any] = None,
+        calibrate: bool = True,
+        calibration_samples: Optional[Sequence[Any]] = None,
+        is_qat: Optional[bool] = False,
+    ):
+        super().__init__(
+            quantizer=quantizer,
+            quantization_config=quantization_config,
+            calibrate=calibrate,
+            calibration_samples=calibration_samples,
+            is_qat=is_qat,
+        )
 
 
 class ToEdgeTransformAndLower(BaseStages.ToEdgeTransformAndLower):
     def __init__(
         self,
         compile_specs: Optional[List[CompileSpec]] = None,
         edge_compile_config: Optional[EdgeCompileConfig] = None,
+        transform_passes: Optional[Union[Sequence[PassType]]] = None,
     ):
         compile_specs = compile_specs or []
         self.partitioners = [EnnPartitioner(compile_specs=compile_specs)]
         self.edge_compile_config = edge_compile_config or get_edge_compile_config()
+        self.transform_passes = transform_passes or get_enn_pass_list()
         self.edge_dialect_program = None
 
     def run(
         self, artifact: ExportedProgram, inputs=None, generate_etrecord: bool = False
     ) -> None:
+        artifact_copy = copy.deepcopy(artifact)
         self.edge_dialect_program = to_edge_transform_and_lower(
-            artifact,
+            artifact_copy,
+            transform_passes=self.transform_passes,
             partitioner=self.partitioners,
             compile_config=self.edge_compile_config,
         )
 
 
 class ToExecutorch(BaseStages.ToExecutorch):
-    pass
+    def run_artifact(self, inputs):
+        runtime_executor = RuntimeExecutor(self.artifact, inputs)
+        return runtime_executor.run_on_device()
 
 
 class SamsungTester(TesterBase):
@@ -60,8 +89,14 @@ def __init__(
     ):
         module.eval()
 
+        stage_classes = TesterBase.default_stage_classes() | {
+            StageType.EXPORT: Export,
+            StageType.TO_EXECUTORCH: ToExecutorch,
+        }
+
         super().__init__(
             module=module,
+            stage_classes=stage_classes,
             example_inputs=example_inputs,
             dynamic_shapes=None,
         )
@@ -71,6 +106,14 @@ def __init__(
         self.example_inputs = example_inputs
         self.compile_specs = compile_specs
 
+    def quantize(self, quantize_stage: Optional[Quantize] = None):
+        if quantize_stage is None:
+            quantizer = EnnQuantizer()
+            quantizer.setup_quant_params(Precision.A8W8)
+            quantize_stage = Quantize(quantizer)
+
+        return super().quantize(quantize_stage)
+
     def to_edge_transform_and_lower(
         self,
         edge_compile_config: Optional[EdgeCompileConfig] = None,

backends/samsung/test/utils/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Samsung Electronics Co. LTD
+# All rights reserved
+#
+# Licensed under the BSD License (the "License"); you may not use this file
+# except in compliance with the License. See the license file in the root
+# directory of this source tree for more details.
+
+from .runtime_executor import RuntimeExecutor
+
+
+__all__ = ["RuntimeExecutor"]

backends/samsung/test/utils/runtime_executor.py

@@ -0,0 +1,163 @@
+import logging
+import os
+import subprocess
+import tempfile
+
+from functools import cache
+from pathlib import Path
+from typing import List, Tuple
+
+import numpy as np
+import torch
+
+
+@cache
+def get_runner_path() -> Path:
+    git_root = subprocess.check_output(
+        ["git", "rev-parse", "--show-toplevel"],
+        cwd=os.path.dirname(os.path.realpath(__file__)),
+        text=True,
+    ).strip()
+    return Path(git_root) / "build_samsung_android/backends/samsung/enn_executor_runner"
+
+
+class ADBTestManager:
+    def __init__(
+        self,
+        pte_file,
+        work_directory,
+        input_files: List[str],
+    ):
+        self.pte_file = pte_file
+        self.work_directory = work_directory
+        self.input_files = input_files
+        self.artifacts_dir = Path(self.pte_file).parent.absolute()
+        self.output_folder = f"{self.work_directory}/output"
+        self.runner = str(get_runner_path())
+
+    def _adb(self, cmd):
+        cmds = ["adb"]
+
+        assert self._is_adb_connected, "Fail to get available device to execute."
+
+        cmds.extend(cmd)
+        command = " ".join(cmds)
+        result = subprocess.run(
+            command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        )
+
+        if result.returncode != 0:
+            logging.info(result.stdout.decode("utf-8").strip())
+            logging.error(result.stderr.decode("utf-8").strip())
+            raise RuntimeError("adb command execute failed")
+
+    def push(self):
+        self._adb(["shell", f"rm -rf {self.work_directory}"])
+        self._adb(["shell", f"mkdir -p {self.work_directory}"])
+        self._adb(["push", self.pte_file, self.work_directory])
+        self._adb(["push", self.runner, self.work_directory])
+
+        for input_file in self.input_files:
+            input_file_path = os.path.join(self.artifacts_dir, input_file)
+            if Path(input_file).name == input_file and os.path.isfile(input_file_path):
+                # default search the same level directory with pte
+                self._adb(["push", input_file_path, self.work_directory])
+            elif os.path.isfile(input_file):
+                self._adb(["push", input_file, self.work_directory])
+            else:
+                raise FileNotFoundError(f"Invalid input file path: {input_file}")
+
+    def execute(self):
+        self._adb(["shell", f"rm -rf {self.output_folder}"])
+        self._adb(["shell", f"mkdir -p {self.output_folder}"])
+        # run the delegation
+        input_files_list = " ".join([os.path.basename(x) for x in self.input_files])
+        enn_executor_runner_args = " ".join(
+            [
+                f"--model {os.path.basename(self.pte_file)}",
+                f'--input "{input_files_list}"',
+                f"--output_path {self.output_folder}",
+            ]
+        )
+        enn_executor_runner_cmd = " ".join(
+            [
+                f"'cd {self.work_directory} &&",
+                f"./enn_executor_runner {enn_executor_runner_args}'",
+            ]
+        )
+
+        self._adb(["shell", f"{enn_executor_runner_cmd}"])
+
+    def pull(self, output_path):
+        self._adb(["pull", "-a", self.output_folder, output_path])
+
+    @staticmethod
+    def _is_adb_connected():
+        try:
+            output = subprocess.check_output(["adb", "devices"])
+            devices = output.decode("utf-8").splitlines()[1:]
+            return [device.split()[0] for device in devices if device.strip() != ""]
+        except subprocess.CAlledProcessError:
+            return False
+
+
+class RuntimeExecutor:
+    def __init__(self, executorch_program, inputs):
+        self.executorch_program = executorch_program
+        self.inputs = inputs
+
+    def run_on_device(self) -> Tuple[torch.Tensor]:
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            pte_filename, input_files = self._save_model_and_inputs(tmp_dir)
+            test_manager = ADBTestManager(
+                pte_file=os.path.join(tmp_dir, pte_filename),
+                work_directory="/data/local/tmp/enn-executorch-test",
+                input_files=input_files,
+            )
+            test_manager.push()
+            test_manager.execute()
+            host_output_save_dir = os.path.join(tmp_dir, "output")
+            test_manager.pull(host_output_save_dir)
+
+            model_outputs = self._get_model_outputs()
+            num_of_output_files = len(os.listdir(host_output_save_dir))
+            assert num_of_output_files == len(
+                model_outputs
+            ), f"Number of outputs is invalid, expect {len(model_outputs)} while got {num_of_output_files}"
+
+            result = []
+            for idx in range(num_of_output_files):
+                output_array = np.fromfile(
+                    os.path.join(host_output_save_dir, f"output_{idx}.bin"),
+                    dtype=np.uint8,
+                )
+                output_tensor = (
+                    torch.from_numpy(output_array)
+                    .view(dtype=model_outputs[idx].dtype)
+                    .view(*model_outputs[idx].shape)
+                )
+                result.append(output_tensor)
+
+            return tuple(result)
+
+    def _get_model_outputs(self):
+        output_node = self.executorch_program.exported_program().graph.output_node()
+        output_fake_tensors = []
+        for ori_output in output_node.args[0]:
+            output_fake_tensors.append(ori_output.meta["val"])
+
+        return tuple(output_fake_tensors)
+
+    def _save_model_and_inputs(self, save_dir):
+        pte_file_name = "program.pte"
+        file_path = os.path.join(save_dir, f"{pte_file_name}")
+        with open(file_path, "wb") as file:
+            self.executorch_program.write_to_file(file)
+
+        inputs_files = []
+        for idx, input in enumerate(self.inputs):
+            input_file_name = f"input_{idx}.bin"
+            input.detach().numpy().tofile(os.path.join(save_dir, input_file_name))
+            inputs_files.append(input_file_name)
+
+        return pte_file_name, inputs_files