NXP backend: Automatic application of RemoveIOQuantOpsPass in NeutronEdgePassManager

Author: StrycekSimon

backends/nxp/edge_passes/neutron_edge_pass_manager.py

@@ -5,6 +5,10 @@
 
 import copy
 
+from executorch.backends.nxp.backend.ir.edge_passes.remove_io_quant_ops_pass import (
+    RemoveIOQuantOpsPass,
+)
+
 from executorch.backends.nxp.edge_passes.move_auxiliary_operator_into_separate_qdq_cluster_pass import (
     MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass,
     MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass,
@@ -24,7 +28,9 @@
 
 class NeutronEdgePassManager(PassManager):
 
-    def __init__(self, passes: list[NeutronEdgePass] = None):
+    def __init__(
+        self, passes: list[NeutronEdgePass] = None, remove_io_quant_ops: bool = False
+    ):
         passes: list[NeutronEdgePass] = passes or [
             MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(),
             MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(),
@@ -35,6 +41,8 @@ def __init__(self, passes: list[NeutronEdgePass] = None):
             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)
@@ -78,12 +86,17 @@ def __call__(self, epm: EdgeProgramManager) -> EdgeProgramManager:
 
             new_programs[name] = new_program
 
-        if len(new_programs) == 0:
-            # No passes were run, return the old EdgeProgramManager.
-            return epm
+        result = epm
 
-        else:
-            # Return a new EdgeProgramManager with the updated programs.
-            return EdgeProgramManager(
+        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,9 +15,6 @@
 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
@@ -115,12 +112,9 @@ def to_quantized_edge_program(
         edge_compile_config=edge_compile_config,
     )
 
-    edge_program_manager = NeutronEdgePassManager()(edge_program_manager)
-
-    if remove_quant_io_ops:
-        edge_program_manager = edge_program_manager.transform(
-            [RemoveIOQuantOpsPass(edge_program_manager=edge_program_manager)]
-        )
+    edge_program_manager = NeutronEdgePassManager(
+        remove_io_quant_ops=remove_quant_io_ops
+    )(edge_program_manager)
 
     compile_spec = generate_neutron_compile_spec(
         target,

examples/nxp/aot_neutron_compile.py

@@ -15,9 +15,6 @@
 import executorch.kernels.quantized  # noqa F401
 
 import torch
-from executorch.backends.nxp.edge_passes.remove_io_quant_ops_pass import (
-    RemoveIOQuantOpsPass,
-)
 from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import (
     NeutronEdgePassManager,
 )
@@ -262,12 +259,9 @@ def _get_batch_size(data):
 
     logging.debug(f"Exported graph:\n{edge_program_manager.exported_program().graph}")
 
-    edge_program_manager = NeutronEdgePassManager()(edge_program_manager)
-
-    if args.remove_quant_io_ops:
-        edge_program_manager = edge_program_manager.transform(
-            [RemoveIOQuantOpsPass(edge_program_manager=edge_program_manager)]
-        )
+    edge_program_manager = NeutronEdgePassManager(
+        remove_io_quant_ops=args.remove_quant_io_ops
+    )(edge_program_manager)
 
     # 5. Delegate to Neutron
     if args.delegate: