Refactor XNNPACK tester to extract base class

Author: GregoryComer

backends/arm/test/misc/test_lifted_tensor.py

@@ -12,6 +12,7 @@
     TosaPipelineBI,
     TosaPipelineMI,
 )
+from executorch.backends.test.harness.stages import StageType
 from executorch.backends.xnnpack.test.tester import ToEdge
 
 
@@ -72,9 +73,8 @@ def test_partition_lifted_tensor_tosa_MI(test_data: input_t1):
         use_to_edge_transform_and_lower=False,
     )
     pipeline.run()
-    to_edge_stage_name = pipeline.tester.stage_name(ToEdge)
     signature = (
-        pipeline.tester.stages[to_edge_stage_name]
+        pipeline.tester.stages[StageType.TO_EDGE]
         .artifact.exported_program()
         .graph_signature
     )
@@ -94,9 +94,8 @@ def test_partition_lifted_tensor_tosa_BI(test_data: input_t1):
         use_to_edge_transform_and_lower=False,
     )
     pipeline.run()
-    to_edge_stage_name = pipeline.tester.stage_name(ToEdge)
     signature = (
-        pipeline.tester.stages[to_edge_stage_name]
+        pipeline.tester.stages[StageType.TO_EDGE]
         .artifact.exported_program()
         .graph_signature
     )

backends/arm/test/tester/analyze_output_utils.py

@@ -14,6 +14,7 @@
     get_output_quantization_params,
 )
 
+from executorch.backends.test.harness.stages import StageType
 from executorch.backends.xnnpack.test.tester.tester import Export, Quantize
 
 logger = logging.getLogger(__name__)
@@ -236,8 +237,8 @@ def dump_error_output(
     if path_to_tosa_files is None:
         path_to_tosa_files = tempfile.mkdtemp(prefix="executorch_result_dump_")
 
-    export_stage = tester.stages.get(tester.stage_name(Export), None)
-    quantize_stage = tester.stages.get(tester.stage_name(Quantize), None)
+    export_stage = tester.stages.get(StageType.EXPORT, None)
+    quantize_stage = tester.stages.get(StageType.QUANTIZE, None)
     if export_stage is not None and quantize_stage is not None:
         output_nodes = get_output_nodes(export_stage.artifact)
         qp_input = get_input_quantization_params(export_stage.artifact)

backends/arm/test/tester/arm_tester.py

@@ -50,6 +50,7 @@
 from executorch.backends.arm.tosa_partitioner import TOSAPartitioner
 from executorch.backends.arm.tosa_specification import TosaSpecification
 
+from executorch.backends.test.harness.stages import StageType
 from executorch.backends.xnnpack.test.tester import Tester
 from executorch.devtools.backend_debug import get_delegation_info
 
@@ -284,13 +285,13 @@ def __init__(
         self.constant_methods = constant_methods
         self.compile_spec = compile_spec
         super().__init__(model, example_inputs, dynamic_shapes)
-        self.pipeline[self.stage_name(InitialModel)] = [
-            self.stage_name(tester.Quantize),
-            self.stage_name(tester.Export),
+        self.pipeline[StageType.INITIAL_MODEL] = [
+            StageType.QUANTIZE,
+            StageType.EXPORT,
         ]
 
         # Initial model needs to be set as a *possible* but not yet added Stage, therefore add None entry.
-        self.stages[self.stage_name(InitialModel)] = None
+        self.stages[StageType.INTIIAL_MODEL] = None
         self._run_stage(InitialModel(self.original_module))
 
     def quantize(self, quantize_stage: Optional[tester.Quantize] = None):
@@ -385,7 +386,7 @@ def serialize(
         return super().serialize(serialize_stage)
 
     def is_quantized(self) -> bool:
-        return self.stages[self.stage_name(tester.Quantize)] is not None
+        return self.stages[StageType.QUANTIZE] is not None
 
     def run_method_and_compare_outputs(
         self,
@@ -414,18 +415,16 @@ def run_method_and_compare_outputs(
         """
 
         if not run_eager_mode:
-            edge_stage = self.stages[self.stage_name(tester.ToEdge)]
+            edge_stage = self.stages[StageType.TO_EDGE]
             if edge_stage is None:
-                edge_stage = self.stages[
-                    self.stage_name(tester.ToEdgeTransformAndLower)
-                ]
+                edge_stage = self.stages[StageType.TO_EDGE_TRANSFORM_AND_LOWER]
             assert (
                 edge_stage is not None
             ), "To compare outputs, at least the ToEdge or ToEdgeTransformAndLower stage needs to be run."
         else:
             # Run models in eager mode. We do this when we want to check that the passes
             # are numerically accurate and the exported graph is correct.
-            export_stage = self.stages[self.stage_name(tester.Export)]
+            export_stage = self.stages[StageType.EXPORT]
             assert (
                 export_stage is not None
             ), "To compare outputs in eager mode, the model must be at Export stage"
@@ -435,11 +434,11 @@ def run_method_and_compare_outputs(
         is_quantized = self.is_quantized()
 
         if is_quantized:
-            reference_stage = self.stages[self.stage_name(tester.Quantize)]
+            reference_stage = self.stages[StageType.QUANTIZE]
         else:
-            reference_stage = self.stages[self.stage_name(InitialModel)]
+            reference_stage = self.stages[StageType.INITIAL_MODEL]
 
-        exported_program = self.stages[self.stage_name(tester.Export)].artifact
+        exported_program = self.stages[StageType.EXPORT].artifact
         output_nodes = get_output_nodes(exported_program)
 
         output_qparams = get_output_quantization_params(output_nodes)
@@ -449,7 +448,7 @@ def run_method_and_compare_outputs(
             quantization_scales.append(getattr(output_qparams[node], "scale", None))
 
         logger.info(
-            f"Comparing Stage '{self.stage_name(test_stage)}' with Stage '{self.stage_name(reference_stage)}'"
+            f"Comparing Stage '{test_stage.stage_type()}' with Stage '{reference_stage.stage_type()}'"
         )
 
         # Loop inputs and compare reference stage with the compared stage.
@@ -500,14 +499,12 @@ def get_graph(self, stage: str | None = None) -> Graph:
             stage = self.cur
         artifact = self.get_artifact(stage)
         if (
-            self.cur == self.stage_name(tester.ToEdge)
-            or self.cur == self.stage_name(Partition)
-            or self.cur == self.stage_name(ToEdgeTransformAndLower)
+            self.cur == StageType.TO_EDGE
+            or self.cur == StageType.PARTITION
+            or self.cur == StageType.TO_EDGE_TRANSFORM_AND_LOWER
         ):
             graph = artifact.exported_program().graph
-        elif self.cur == self.stage_name(tester.Export) or self.cur == self.stage_name(
-            tester.Quantize
-        ):
+        elif self.cur == StageType.EXPORT or self.cur == StageType.QUANTIZE:
             graph = artifact.graph
         else:
             raise RuntimeError(
@@ -533,8 +530,8 @@ def dump_operator_distribution(
         if (
             self.cur
             in (
-                self.stage_name(tester.Partition),
-                self.stage_name(ToEdgeTransformAndLower),
+                StageType.PARTITION,
+                StageType.TO_EDGE_TRANSFORM_AND_LOWER,
             )
             and print_table
         ):
@@ -625,7 +622,7 @@ def run_transform_for_annotation_pipeline(
             stage = self.cur
         # We need to clone the artifact in order to ensure that the state_dict is preserved after passes are run.
         artifact = self.get_artifact(stage)
-        if self.cur == self.stage_name(tester.Export):
+        if self.cur == StageType.EXPORT:
             new_gm = ArmPassManager(get_tosa_spec(self.compile_spec)).transform_for_annotation_pipeline(  # type: ignore[arg-type]
                 graph_module=artifact.graph_module
             )

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