Refactor XNNPACK tester to extract base class

Author: GregoryComer

backends/test/harness/__init__.py

@@ -0,0 +1 @@
+from .tester import Tester

backends/test/harness/stages/__init__.py

@@ -0,0 +1,10 @@
+from .export import Export
+from .partition import Partition
+from .quantize import Quantize
+from .run_passes import RunPasses
+from .serialize import Serialize
+from .stage import Stage, StageType
+from .to_edge import ToEdge
+from .to_edge_transform_and_lower import ToEdgeTransformAndLower
+from .to_executorch import ToExecutorch
+

backends/test/harness/stages/export.py

@@ -0,0 +1,32 @@
+from typing import Any, Optional, Sequence, Tuple
+
+import torch
+
+from executorch.backends.test.harness.stages.stage import Stage, StageType
+from torch.export import export, ExportedProgram
+
+class Export(Stage):
+    def __init__(self, dynamic_shapes: Optional[Tuple[Any]] = None):
+        self.exported_program = None
+        self.dynamic_shapes = dynamic_shapes
+
+    def stage_type(self) -> StageType:
+        return StageType.EXPORT
+
+    def run(
+        self,
+        artifact: torch.nn.Module,
+        inputs: Tuple[torch.Tensor],
+    ) -> None:
+        self.exported_program = export(
+            artifact, inputs, dynamic_shapes=self.dynamic_shapes, strict=True
+        )
+
+    @property
+    def artifact(self) -> ExportedProgram:
+        return self.exported_program
+
+    @property
+    def graph_module(self) -> str:
+        return self.exported_program.graph_module
+

backends/test/harness/stages/partition.py

@@ -0,0 +1,27 @@
+from executorch.backends.test.harness.stages.stage import Stage, StageType
+from executorch.exir import (
+    EdgeProgramManager,
+)
+from executorch.exir.backend.backend_api import validation_disabled
+from executorch.exir.backend.partitioner import Partitioner
+
+class Partition(Stage):
+    def __init__(self, partitioner: Partitioner):
+        self.partitioner = partitioner
+        self.delegate_module = None
+    
+    def stage_type(self) -> StageType:
+        return StageType.PARTITION
+
+    def run(self, artifact: EdgeProgramManager, inputs=None):
+        with validation_disabled():
+            self.delegate_module = artifact
+            self.delegate_module = self.delegate_module.to_backend(self.partitioner)
+
+    @property
+    def artifact(self) -> EdgeProgramManager:
+        return self.delegate_module
+
+    @property
+    def graph_module(self) -> str:
+        return self.delegate_module.exported_program().graph_module

backends/test/harness/stages/quantize.py

@@ -0,0 +1,78 @@
+from typing import Any, Optional, Sequence, Tuple
+
+import torch
+
+from executorch.backends.test.harness.stages.stage import Stage, StageType
+from executorch.backends.transforms.duplicate_dynamic_quant_chain import (
+    DuplicateDynamicQuantChainPass,
+)
+
+from torch.export import export_for_training
+
+from torchao.quantization.pt2e.quantize_pt2e import (
+    convert_pt2e,
+    prepare_pt2e,
+    prepare_qat_pt2e,
+)
+from torchao.quantization.pt2e.quantizer import Quantizer
+
+class Quantize(Stage):
+    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,
+    ):
+        self.quantizer = quantizer
+        self.quantization_config = quantization_config
+        self.calibrate = calibrate
+        self.calibration_samples = calibration_samples
+
+        self.quantizer.set_global(self.quantization_config)
+
+        self.converted_graph = None
+        self.is_qat = is_qat
+    
+    def stage_type(self) -> str:
+        return StageType.QUANTIZE
+
+    def run(
+        self, artifact: torch.nn.Module, inputs: Optional[Tuple[torch.Tensor]]
+    ) -> None:
+        assert inputs is not None
+        if self.is_qat:
+            artifact.train()
+        captured_graph = export_for_training(artifact, inputs, strict=True).module()
+
+        assert isinstance(captured_graph, torch.fx.GraphModule)
+
+        if self.is_qat:
+            prepared = prepare_qat_pt2e(captured_graph, self.quantizer)
+        else:
+            prepared = prepare_pt2e(captured_graph, self.quantizer)
+
+        if self.calibrate:
+            # Calibrate prepared model to provide data to quantization observers.
+            if self.calibration_samples is not None:
+                for inp in self.calibration_samples:
+                    prepared(*inp)
+            else:
+                prepared(*inputs)
+
+        converted = convert_pt2e(prepared)
+        DuplicateDynamicQuantChainPass()(converted)
+
+        self.converted_graph = converted
+
+    @property
+    def artifact(self) -> torch.fx.GraphModule:
+        return self.converted_graph
+
+    @property
+    def graph_module(self) -> str:
+        return self.converted_graph
+
+    def run_artifact(self, inputs):
+        return self.converted_graph.forward(*inputs)

backends/test/harness/stages/run_passes.py

@@ -0,0 +1,70 @@
+from typing import Any, Callable, List, Optional, Sequence, Type, Tuple, Union
+
+import torch
+
+from executorch.backends.test.harness.stages.stage import Stage, StageType
+from executorch.exir import (
+    EdgeCompileConfig,
+    EdgeProgramManager,
+)
+from executorch.exir.backend.partitioner import Partitioner
+from torch._export.pass_base import PassType
+from torch.export import ExportedProgram
+
+class RunPasses(Stage):
+    def __init__(
+        self,
+        pass_manager_cls: Type,
+        pass_list: Optional[List[Type[PassType]]] = None,
+        pass_functions: Optional[List[Callable]] = None,
+    ):
+        self.pass_manager_cls = pass_manager_cls
+        self.pass_list = pass_list
+        self.pass_functions = pass_functions
+        self.edge_or_aten_program = None
+    
+    def stage_type(self) -> StageType:
+        return StageType.RUN_PASSES
+
+    def run(
+        self, artifact: Union[EdgeProgramManager, ExportedProgram], inputs=None
+    ) -> None:
+        if isinstance(artifact, EdgeProgramManager):
+            self.edge_or_aten_program = artifact
+            if self.pass_list:
+                pass_manager = self.pass_manager_cls(
+                    artifact.exported_program(), self.pass_list
+                )
+                self.edge_or_aten_program._edge_programs["forward"] = (
+                    pass_manager.transform()
+                )
+            if self.pass_functions:
+                assert isinstance(self.pass_functions, list)
+                for pass_function in self.pass_functions:
+                    self.edge_or_aten_program._edge_programs["forward"] = pass_function(
+                        self.edge_or_aten_program.exported_program()
+                    )
+        else:
+            transformed_ep = artifact
+            if self.pass_list:
+                assert isinstance(self.pass_list, list)
+                for pass_ in self.pass_list:
+                    transformed_ep = _transform(transformed_ep, pass_())
+
+            if self.pass_functions:
+                assert isinstance(self.pass_functions, list)
+                for pass_function in self.pass_functions:
+                    transformed_ep = pass_function(transformed_ep)
+
+            self.edge_or_aten_program = transformed_ep
+
+    @property
+    def artifact(self) -> Union[EdgeProgramManager, ExportedProgram]:
+        return self.edge_or_aten_program
+
+    @property
+    def graph_module(self) -> str:
+        if isinstance(self.edge_or_aten_program, EdgeProgramManager):
+            return self.edge_or_aten_program.exported_program().graph_module
+        else:
+            return self.edge_or_aten_program.graph_module

backends/test/harness/stages/serialize.py

@@ -0,0 +1,60 @@
+import copy
+import logging
+
+from typing import Optional
+
+from executorch.backends.test.harness.stages.stage import Stage, StageType
+from executorch.exir import (
+    EdgeCompileConfig,
+    EdgeProgramManager,
+    ExecutorchBackendConfig,
+    ExecutorchProgramManager,
+)
+
+from torch.utils._pytree import tree_flatten
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+try:
+    from executorch.extension.pybindings.portable_lib import (  # @manual
+        _load_for_executorch_from_buffer,
+    )
+except ImportError as e:
+    logger.warning(f"{e=}")
+    pass
+
+class Serialize(Stage):
+    def __init__(self):
+        self.buffer = None
+    
+    def stage_type(self) -> StageType:
+        return StageType.SERIALIZE
+
+    def run(self, artifact: ExecutorchProgramManager, inputs=None) -> None:
+        self.buffer = artifact.buffer
+
+    @property
+    def artifact(self) -> bytes:
+        return self.buffer
+
+    @property
+    def graph_module(self) -> None:
+        return None
+
+    def run_artifact(self, inputs):
+        inputs_flattened, _ = tree_flatten(inputs)
+        executorch_module = _load_for_executorch_from_buffer(self.buffer)
+        executorch_output = copy.deepcopy(
+            executorch_module.run_method("forward", tuple(inputs_flattened))
+        )
+        return executorch_output
+
+    def dump_artifact(self, path_to_dump: Optional[str]):
+        """
+        dump_artifact is overridden to dump the serialized bytes into pte file
+        """
+        if not path_to_dump:
+            raise RuntimeError("path_to_dump file not provided")
+        else:
+            with open(path_to_dump, "wb") as f:
+                f.write(self.artifact)

backends/test/harness/stages/stage.py

@@ -0,0 +1,90 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Optional
+
+from executorch.exir import (
+    EdgeProgramManager,
+)
+
+
+class StageType(Enum):
+    QUANTIZE = 0
+    EXPORT = 1
+    RUN_PASSES = 2
+    TO_EDGE = 3
+    TO_EDGE_TRANSFORM_AND_LOWER = 4
+    PARTITION = 5
+    TO_EXECUTORCH = 6
+    SERIALIZE = 7
+
+
+class Stage(ABC):
+    """
+    Interface for a Stage in the PT2.0 lowering pipeline
+    """
+
+    @abstractmethod
+    def stage_type(self) -> StageType:
+        """
+        Returns the type of the stage.
+        """
+        pass
+
+    @abstractmethod
+    def run(self, artifact, inputs):
+        """
+        Executes this stage, generates the 'artifact', for later stages.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def artifact(self):
+        """
+        Returns the artifact generated by this stage. To be used by the next stage in the pipeline.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def graph_module(self):
+        """
+        Return the artifact's graph module for this stage
+        """
+        pass
+
+    def run_artifact(self, inputs):
+        """
+        Returns the output of calling the artifact generated by this stage with inputs
+        """
+        if isinstance(self.artifact, ExportedProgram):
+            return self.artifact(*inputs)
+        else:
+            return self.artifact.exported_program().module()(*inputs)
+
+    # Debug Tools for stages
+    def artifact_str(self):
+        """
+        Return string printable artifact for this stage
+        """
+        if isinstance(self.artifact, EdgeProgramManager):
+            return self.artifact.exported_program()
+        return self.artifact
+
+    def stage_banner(self):
+        """
+        Returns banner string for this stage
+        """
+        return "#" * 36 + " " + str(self.__class__.__name__) + " " + "#" * 36 + "\n"
+
+    def dump_artifact(self, path_to_dump: Optional[str]):
+        """
+        Dumps string printable artifact to path. If path_to_dump, then it is printed to terminal
+        """
+        if path_to_dump:
+            with open(path_to_dump, "a") as fp:
+                fp.write(str(self.stage_banner() + "\n"))
+                fp.write(str(self.artifact_str()))
+        else:
+            print(self.stage_banner() + "\n")
+            print(self.artifact_str())