Unify pipeline workflow, use to_edge_transform_and_lower() in executorch_pipeline.py, fix edge pass manager usage

Author: roman-janik-nxp

backends/nxp/edge_passes/neutron_edge_pass_manager.py

@@ -3,34 +3,17 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import copy
-
 from executorch.backends.nxp.edge_passes.move_auxiliary_operator_into_separate_qdq_cluster_pass import (
     MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass,
     MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass,
 )
 from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass
-
-from executorch.backends.nxp.edge_passes.remove_io_quant_ops_pass import (
-    RemoveIOQuantOpsPass,
-)
-from executorch.exir import EdgeProgramManager
-from executorch.exir.program._program import (
-    _get_updated_graph_signature,
-    _get_updated_range_constraints,
-)
-
-from torch import nn
-from torch.export import ExportedProgram
-from torch.fx.passes.infra.pass_base import PassResult
 from torch.fx.passes.infra.pass_manager import PassManager
 
 
 class NeutronEdgePassManager(PassManager):
 
-    def __init__(
-        self, passes: list[NeutronEdgePass] = None, remove_io_quant_ops: bool = False
-    ):
+    def __init__(self, passes: list[NeutronEdgePass] = None):
         passes: list[NeutronEdgePass] = passes or [
             MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(),
             MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(),
@@ -40,63 +23,3 @@ def __init__(
             passes,
             steps=10,  # Empirical value. At most 10 cycles of passes will be run.
         )
-
-        self.remove_io_quant_ops = remove_io_quant_ops
-
-    def _transform_graph_module(self, module: nn.Module) -> PassResult:
-        """Apply the passes to a single graph module."""
-        pass_result: PassResult = super().__call__(module)
-
-        graph_module = pass_result.graph_module
-        graph_module.graph.eliminate_dead_code()
-        graph_module.recompile()
-
-        return pass_result
-
-    def __call__(self, epm: EdgeProgramManager) -> EdgeProgramManager:
-        """Apply the passes to all graph modules in the edge program."""
-        new_programs: dict[str, ExportedProgram] = {}
-
-        for name, program in epm._edge_programs.items():
-            pass_result = self._transform_graph_module(program.graph_module)
-
-            if pass_result.modified:
-                # Create a new exported program.
-                new_program = ExportedProgram(
-                    root=pass_result.graph_module,
-                    graph=pass_result.graph_module.graph,
-                    graph_signature=_get_updated_graph_signature(
-                        program.graph_signature, pass_result.graph_module
-                    ),
-                    state_dict=program.state_dict,
-                    range_constraints=_get_updated_range_constraints(
-                        pass_result.graph_module
-                    ),
-                    module_call_graph=copy.deepcopy(program._module_call_graph),
-                    example_inputs=program.example_inputs,
-                    constants=program.constants,
-                    verifiers=[program.verifier],
-                )
-                new_program.graph_module.meta.update(program.graph_module.meta)
-                new_program.graph_module.meta.update(pass_result.graph_module.meta)
-
-            else:
-                # Keep the old exported program.
-                new_program = program
-
-            new_programs[name] = new_program
-
-        result = epm
-
-        if len(new_programs) > 0:
-            # Use a new EdgeProgramManager with the updated programs if any update was performed.
-            result = EdgeProgramManager(
-                new_programs, copy.deepcopy(epm._config_methods), epm.compile_config
-            )
-
-        if self.remove_io_quant_ops:
-            result = result.transform(
-                [RemoveIOQuantOpsPass(edge_program_manager=result)]
-            )
-
-        return result

backends/nxp/tests/executorch_pipeline.py

@@ -15,6 +15,9 @@
 from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import (
     NeutronEdgePassManager,
 )
+from executorch.backends.nxp.edge_passes.remove_io_quant_ops_pass import (
+    RemoveIOQuantOpsPass,
+)
 from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
 from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec
 from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
@@ -23,9 +26,10 @@
     EdgeProgramManager,
     ExecutorchBackendConfig,
     ExecutorchProgramManager,
+    to_edge_transform_and_lower,
 )
-from executorch.extension.export_util.utils import export_to_edge
 from torch import nn
+from torch.export import export
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
 
@@ -105,24 +109,24 @@ def to_quantized_edge_program(
         calibration_inputs,
     )
 
-    edge_compile_config = EdgeCompileConfig(_check_ir_validity=False)
-    edge_program_manager = export_to_edge(
-        exir_program_aten__module_quant,
-        example_input,
-        edge_compile_config=edge_compile_config,
-    )
-
-    edge_program_manager = NeutronEdgePassManager(
-        remove_io_quant_ops=remove_quant_io_ops
-    )(edge_program_manager)
-
     compile_spec = generate_neutron_compile_spec(
         target,
         operators_not_to_delegate=operators_not_to_delegate,
         neutron_converter_flavor=neutron_converter_flavor,
     )
-    partitioner = NeutronPartitioner(compile_spec, custom_delegation_options)
-    edge_program_manager = edge_program_manager.to_backend(partitioner)
+    partitioners = [NeutronPartitioner(compile_spec, custom_delegation_options)]
+
+    edge_program_manager = to_edge_transform_and_lower(
+        export(exir_program_aten__module_quant, example_input, strict=True),
+        transform_passes=NeutronEdgePassManager(),
+        partitioner=partitioners,
+        compile_config=EdgeCompileConfig(_check_ir_validity=False),
+    )
+
+    if remove_quant_io_ops:
+        edge_program_manager = edge_program_manager.transform(
+            [RemoveIOQuantOpsPass(edge_program_manager=edge_program_manager)]
+        )
 
     return edge_program_manager
 

examples/nxp/aot_neutron_compile.py

@@ -18,6 +18,9 @@
 from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import (
     NeutronEdgePassManager,
 )
+from executorch.backends.nxp.edge_passes.remove_io_quant_ops_pass import (
+    RemoveIOQuantOpsPass,
+)
 from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
 from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec
 from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
@@ -273,13 +276,15 @@ def _get_batch_size(data):
 
     edge_program_manager = to_edge_transform_and_lower(
         export(module, example_inputs, strict=True),
+        transform_passes=NeutronEdgePassManager(),
         partitioner=partitioners,
         compile_config=EdgeCompileConfig(),
     )
 
-    edge_program_manager = NeutronEdgePassManager(
-        remove_io_quant_ops=args.remove_quant_io_ops
-    )(edge_program_manager)
+    if args.remove_quant_io_ops:
+        edge_program_manager = edge_program_manager.transform(
+            [RemoveIOQuantOpsPass(edge_program_manager=edge_program_manager)]
+        )
 
     logging.debug(f"Lowered graph:\n{edge_program_manager.exported_program().graph}")