Pull request pytorch#77: [EIEX-187] Add option to specify neutron-converter flavor

Merge in AITEC/executorch from feature/nxf93343/EIEX-187-neutron-converter-flavor-specification to main-nxp

* commit 'ef00ebf93d241b175461c6a3c30d799ee138cce6':
  Add option to specify neutron-converter flavor

Author: skywall

backends/nxp/backend/neutron_converter_manager.py

@@ -3,8 +3,10 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+import importlib
+import pkgutil
 
-from neutron_converter_wrapper import neutron_converter
+from executorch.backends.nxp.backend.ir.converter.node_converter import Target
 
 
 class NeutronConverterManager:
@@ -13,7 +15,27 @@ class NeutronConverterManager:
     contains NeutronGraph nodes.
     """
 
-    def convert(self, tflite_model: bytes, target: str) -> bytes:
+    _supported_target_names =  [Target.RT700.value]
+
+    def convert(self, tflite_model: bytes, target: str, neutron_converter_flavor: str) -> bytes:
+        # Neutron converter crashes if we provide invalid target -> verify.
+        if target not in self._supported_target_names:
+            raise RuntimeError(f"Target '{target}' is not supported by NeutronConverterManager.")
+
+        neutron_converter_modules = [module.name for module in pkgutil.iter_modules() if
+                                     module.name.startswith("neutron_converter")]
+
+        requested_module_name = f"neutron_converter_{neutron_converter_flavor}"
+        if requested_module_name not in neutron_converter_modules:
+            if len(neutron_converter_modules) > 0:
+                raise RuntimeError(f"Neutron Converter module with flavor '{neutron_converter_flavor}' "
+                                   f"not found. Available modules: {neutron_converter_modules}.")
+            else:
+                raise RuntimeError(f"Neutron Converter module with flavor '{neutron_converter_flavor}' "
+                                   f"not found. Install 'neutron_converter_[flavor]' Python package.")
+
+        neutron_converter = importlib.import_module(f"{requested_module_name}.neutron_converter")
+
         cctx = neutron_converter.CompilationContext()
         cctx.targetOpts = neutron_converter.getNeutronTarget(target)
         model_converted = neutron_converter.convertModel(list(tflite_model), cctx)

backends/nxp/nxp_backend.py

@@ -35,6 +35,7 @@ def __init__(self):
         self.compiler_flags = []
         self.output_format = None
         self.operators_not_to_delegate: List[str] = []
+        self.neutron_converter_flavor = None
 
     def _replace_colons(self, operator: str) -> str:
         """
@@ -45,6 +46,7 @@ def _replace_colons(self, operator: str) -> str:
     def neutron_compile_spec(
         self,
         config: str,
+        neutron_converter_flavor: str,
         extra_flags: Optional[str] = None,
         operators_not_to_delegate: Optional[List[str]] = None,
     ):
@@ -61,16 +63,15 @@ def neutron_compile_spec(
         except ValueError:
             raise ValueError(f'Config `{config}` is not a valid target. Must be one of `{Target.values()}`.')
 
-        assert (
-            self.output_format is None
-        ), f"Output format already set to f{self.output_format}"
+        self.neutron_converter_flavor = neutron_converter_flavor
+
+        assert (self.output_format is None), f"Output format already set to f{self.output_format}"
         self.output_format = "tflite"
-        self.compiler_flags = [
+        self.compiler_flags = []
 
-        ]
         if extra_flags is not None:
             self.compiler_flags.append(extra_flags)
-            
+
         if operators_not_to_delegate is not None:
             self.operators_not_to_delegate = [self._replace_colons(op) for op in operators_not_to_delegate]
 
@@ -85,6 +86,7 @@ def build(self):
                 CompileSpec("output_format", "tflite".encode()),
                 CompileSpec("compile_flags", " ".join(self.compiler_flags).encode()),
                 CompileSpec("target", self.config.value.encode()),
+                CompileSpec("neutron_converter_flavor", self.neutron_converter_flavor.encode()),
                 CompileSpec("operators_not_to_delegate", ",".join(self.operators_not_to_delegate).encode())
             ]
 
@@ -93,12 +95,14 @@ def build(self):
 
 def generate_neutron_compile_spec(
     config: str,  # The target platform. For example "imxrt700".
+    neutron_converter_flavor: str,
     system_config: Optional[str] = None,
     extra_flags: Optional[str] = None,
     operators_not_to_delegate: Optional[List[str]] = None,
 ) -> List[CompileSpec]:
     return NeutronCompileSpecBuilder().neutron_compile_spec(
         config,
+        neutron_converter_flavor,
         extra_flags=extra_flags,
         operators_not_to_delegate=operators_not_to_delegate
     ).build()
@@ -123,6 +127,7 @@ def preprocess(
         compile_flags = []
         binary = bytes()
         target = ""
+        neutron_converter_flavor = ""
         for spec in compile_spec:
             if spec.key == "output_format":
                 output_format = spec.value.decode()
@@ -132,6 +137,8 @@ def preprocess(
                 compile_flags.append(spec.value.decode())
             if spec.key == "operators_not_to_delegate":
                 operators_not_to_delegate = spec.value.decode().split(',')
+            if spec.key == "neutron_converter_flavor":
+                neutron_converter_flavor = spec.value.decode()
 
         # Check that the output format is set in the compile spec
         if not output_format:
@@ -156,7 +163,7 @@ def preprocess(
             # Convert the edge program to TFLite.
             tflite_model, io_formats = EdgeProgramToIRConverter().convert_program(edge_program)
 
-            neutron_model = NeutronConverterManager().convert(tflite_model, target)
+            neutron_model = NeutronConverterManager().convert(tflite_model, target, neutron_converter_flavor)
 
             # Dump the tflite file if logging level is enabled
             if logging.root.isEnabledFor(logging.DEBUG):
@@ -169,7 +176,6 @@ def preprocess(
                 NeutronBackend.counter = NeutronBackend.counter + 1
 
             binary = PayloadComposer().get_binary_payload(io_formats, neutron_model)
-
         else:
             raise RuntimeError(f"Unknown format {output_format}")
 

backends/nxp/tests/executorch_pipeline.py

@@ -27,7 +27,8 @@ def _quantize_model(model, calibration_inputs: list[tuple[torch.Tensor]]):
     return m
 
 
-def to_quantized_edge_program(model: torch.nn.Module, input_shape: tuple, operators_not_to_delegate: list[str] = None, target="imxrt700") -> EdgeProgramManager:
+def to_quantized_edge_program(model: torch.nn.Module, input_shape: tuple, operators_not_to_delegate: list[str] = None,
+                              target="imxrt700", neutron_converter_flavor="wrapper") -> EdgeProgramManager:
     calibration_inputs = [(torch.randn(input_shape),), (torch.randn(input_shape),)]
     example_input = (torch.ones(*input_shape),)
 
@@ -40,7 +41,8 @@ def to_quantized_edge_program(model: torch.nn.Module, input_shape: tuple, operat
     exir_program_aten_quant = _quantize_model(exir_program_aten, calibration_inputs)
     edge_program_manager = export_to_edge(exir_program_aten_quant, example_input)
 
-    compile_spec = generate_neutron_compile_spec(target, operators_not_to_delegate=operators_not_to_delegate) if operators_not_to_delegate else generate_neutron_compile_spec(target)
+    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 = edge_program_manager.to_backend(partitioner)

backends/nxp/tests/test_neutron_backend.py

@@ -37,7 +37,7 @@ def test_neutron_backend__single_conv_model__payload_header_channels_last():
 
 
 def test_neutron_backend__linear_softmax_model__payload_header_formatless():
-    edge_program_manager = to_quantized_edge_program(LinearSoftmaxModule(), (1, 12), target=Target.IGNORE)
+    edge_program_manager = to_quantized_edge_program(LinearSoftmaxModule(), (1, 12))
     payload = edge_program_manager.exported_program().graph_module.lowered_module_0.processed_bytes
 
     assert payload[0] == 0x1  # Single input
@@ -85,14 +85,14 @@ def test_lowered_program_and_tflite_output_match__conv2d__no_bias(mocker):
     assert np.max(np.abs(output_edge - output_tflite)) <= 1
 
 
-def test_conv_fc_softmax__lowered_program_and_tflite_output_match(mocker):
+def test_conv_fc__lowered_program_and_tflite_output_match(mocker):
     converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
 
     model = ConvFCSoftmaxModule()
     input_shape = (1, 4, 5, 5)
 
     # Run conversion
-    _ = to_quantized_edge_program(model, input_shape, target=Target.IGNORE)
+    _ = to_quantized_edge_program(model, input_shape)
 
     # Capture converted program
     exported_program: ExportedProgram = converter_spy.call_args.args[1]
@@ -103,11 +103,10 @@ def test_conv_fc_softmax__lowered_program_and_tflite_output_match(mocker):
     # No Transpose ops in produced TFLite model
     tflite_subgraph = Model.GetRootAs(tflite_flatbuffers_model).Subgraphs(0)
 
-    assert tflite_subgraph.OperatorsLength() == 4
+    assert tflite_subgraph.OperatorsLength() == 3
     assert tflite_subgraph.Operators(0).BuiltinOptionsType() == BuiltinOptions.Conv2DOptions
     assert tflite_subgraph.Operators(1).BuiltinOptionsType() == BuiltinOptions.ReshapeOptions
     assert tflite_subgraph.Operators(2).BuiltinOptionsType() == BuiltinOptions.FullyConnectedOptions
-    assert tflite_subgraph.Operators(3).BuiltinOptionsType() == BuiltinOptions.SoftmaxOptions
 
     # Verify outputs of program and TFLite model
     input_data = (torch.randn(input_shape, dtype=torch.float32)).type(torch.int8).detach().numpy()

backends/nxp/tests/test_neutron_converter_manager.py

@@ -7,9 +7,7 @@
 from executorch.backends.nxp.tests.models import Conv2dModule
 
 
-def test_conv2d_neutron_conversion():
-    pytest.importorskip("neutron_converter_wrapper")
-
+def test_conv2d_neutron_conversion__default_flavor():
     model = Conv2dModule()
 
     example_input = (torch.ones(1, 4, 32, 32),)
@@ -20,6 +18,23 @@ def test_conv2d_neutron_conversion():
     tflite_model, _ = edge_program_converter.convert_program(edge_program_manager.exported_program())
 
     neutron_converter_manager = NeutronConverterManager()
-    neutron_model = neutron_converter_manager.convert(tflite_model, "imxrt700")
+    neutron_model = neutron_converter_manager.convert(tflite_model, "imxrt700", "wrapper")
 
     assert len(neutron_model), "Produced NeutronGraph-based TFLite model has zero length!"
+
+
+def test__conv2d_neutron_conversion__invalid_flavor():
+    model = Conv2dModule()
+
+    example_input = (torch.ones(1, 4, 32, 32),)
+    exir_program = torch.export.export(model, example_input)
+    edge_program_manager = exir.to_edge(exir_program)
+
+    edge_program_converter = EdgeProgramToIRConverter()
+    tflite_model, _ = edge_program_converter.convert_program(edge_program_manager.exported_program())
+
+    neutron_converter_manager = NeutronConverterManager()
+    with pytest.raises(RuntimeError) as excinfo:
+        _ = neutron_converter_manager.convert(tflite_model, "imxrt700", "bad_flavor")
+
+    assert "Neutron Converter module with flavor 'bad_flavor' not found." in str(excinfo)

backends/nxp/tests/test_operator_selector.py

@@ -10,7 +10,7 @@ def test_operator_selector_mechanism():
 
     operators_not_to_delegate = ["aten::convolution"]
 
-    edge_program_manager = to_quantized_edge_program(model, input_shape, target=Target.IGNORE, operators_not_to_delegate=operators_not_to_delegate)
+    edge_program_manager = to_quantized_edge_program(model, input_shape, operators_not_to_delegate=operators_not_to_delegate)
 
     exported_program = edge_program_manager.exported_program()
 

examples/nxp/aot_neutron_compile.py

@@ -142,6 +142,13 @@ def post_training_quantize(model, calibration_inputs: tuple[torch.Tensor] | Iter
         default="imxrt700",
         help="Platform for running the delegated model",
     )
+    parser.add_argument(
+        "-c", "--neutron_converter_flavor",
+        required=False,
+        default="wrapper",
+        help="Flavor of installed neutron-converter module. Neutron-converter module named "
+             "'neutron_converter_SDK_24_12' has flavor 'SDK_24_12'.",
+    )
     parser.add_argument(
         "-q",
         "--quantize",
@@ -246,6 +253,7 @@ def post_training_quantize(model, calibration_inputs: tuple[torch.Tensor] | Iter
             NeutronPartitioner(
                 generate_neutron_compile_spec(
                     args.target,
+                    args.neutron_converter_flavor,
                     operators_not_to_delegate=args.operators_not_to_delegate
                 )
             )