NXP backend: Add infrastructure for running pre-processing passes on edge dialect programs.

Author: MartinPavella

backends/nxp/edge_passes/neutron_edge_pass.py

@@ -0,0 +1,54 @@
+# Copyright 2025 NXP
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from abc import abstractmethod
+
+import torch
+from torch.fx.passes.infra.pass_base import PassResult
+
+from executorch.exir.pass_base import ExportPass
+
+
+class NeutronEdgePass(ExportPass):
+    """ Abstract parent class for pre-processing passes on the edge dialect level. """
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        """ Call `self.run()` as long as changes are being made. After a pass modifies the graph, it cannot keep on
+             iterating through its nodes, and must return. This method allows the pass to go through the whole model.
+        """
+
+        # Every pass will return once it makes a change to the graph, to avoid traversing and modifying a graph at the
+        #  same time. Therefore, it must be called multiple times (at most `iteration_limit` times).
+        iteration_limit = len(graph_module.graph.nodes)
+        modified = False
+        for _ in range(iteration_limit):
+            res = self.run(graph_module)
+            if res.modified:
+                modified = True
+                graph_module = res.graph_module
+
+            else:
+                # No more changes have been made.
+                graph_module = self.recompile_module(graph_module)
+                return PassResult(graph_module, modified)
+
+        # Iteration limit was reached.
+        logging.warning(f'The NeutronEdgePass `{self.__class__.__name__}` reached the iteration limit.')
+        graph_module = self.recompile_module(graph_module)
+        return PassResult(graph_module, modified)
+
+    @abstractmethod
+    def run(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        """ Child classes should implement their graph modification here. """
+        pass
+
+    def recompile_module(
+        self, graph_module: torch.fx.GraphModule
+    ) -> torch.fx.GraphModule:
+        """ Recompile the graph and re-trace the metadata. This should ensure that the datatypes and shapes are correct.
+        """
+        graph_module.recompile()
+        return super().call(graph_module).graph_module

backends/nxp/edge_passes/neutron_edge_pass_manager.py

@@ -0,0 +1,80 @@
+# Copyright 2025 NXP
+#
+# 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 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
+
+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.exir import EdgeProgramManager
+from executorch.exir.program._program import _get_updated_graph_signature, _get_updated_range_constraints
+
+
+class NeutronEdgePassManager(PassManager):
+
+    def __init__(self, passes: list[NeutronEdgePass] = None):
+        passes: list[NeutronEdgePass] = passes or [
+        ]
+
+        super().__init__(
+            passes,
+            steps=10  # Empirical value. At most 10 cycles of passes will be run.
+        )
+
+    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
+
+        if len(new_programs) == 0:
+            # No passes were run, return the old EdgeProgramManager.
+            return epm
+
+        else:
+            # Return a new EdgeProgramManager with the updated programs.
+            return EdgeProgramManager(new_programs, copy.deepcopy(epm._config_methods), epm.compile_config)

backends/nxp/tests/executorch_pipeline.py

@@ -4,8 +4,11 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
+from torch import nn
+from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
 from executorch import exir
+from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import NeutronEdgePassManager
 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
@@ -14,10 +17,8 @@
     EdgeProgramManager,
     ExecutorchBackendConfig,
     ExecutorchProgramManager,
-    to_edge_transform_and_lower,
 )
-from torch import nn
-from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
+from executorch.extension.export_util.utils import export_to_edge
 
 
 def _quantize_model(model, calibration_inputs: list[tuple[torch.Tensor]]):
@@ -49,19 +50,22 @@ def to_quantized_edge_program(
         exir_program_aten.module(), 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()(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)
-    edge_program_manager = to_edge_transform_and_lower(
-        torch.export.export(
-            exir_program_aten__module_quant, example_input, strict=True
-        ),
-        partitioner=[partitioner],
-        compile_config=EdgeCompileConfig(_check_ir_validity=False),
-    )
+    edge_program_manager = edge_program_manager.to_backend(partitioner)
 
     return edge_program_manager