From 4c575c7eb015817280b263dc79bc47c4e547601d Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Fri, 31 Jul 2026 14:56:29 +0200 Subject: [PATCH] NXB backend: Add recipes for Neutron backend lowering. --- .../edge_passes/neutron_edge_pass_manager.py | 2 +- backends/nxp/recipes/nxp_recipe_provider.py | 353 ++++++++++++++++ backends/nxp/recipes/nxp_recipe_types.py | 35 ++ backends/nxp/tests/executorch_pipeline.py | 9 + .../tests/generic_tests/test_recipe_export.py | 387 ++++++++++++++++++ export/recipe.py | 80 +++- export/stages.py | 56 ++- 7 files changed, 909 insertions(+), 13 deletions(-) create mode 100644 backends/nxp/recipes/nxp_recipe_provider.py create mode 100644 backends/nxp/recipes/nxp_recipe_types.py create mode 100644 backends/nxp/tests/generic_tests/test_recipe_export.py diff --git a/backends/nxp/edge_passes/neutron_edge_pass_manager.py b/backends/nxp/edge_passes/neutron_edge_pass_manager.py index 3a7fc3ffbfa..e95f8d18144 100644 --- a/backends/nxp/edge_passes/neutron_edge_pass_manager.py +++ b/backends/nxp/edge_passes/neutron_edge_pass_manager.py @@ -17,7 +17,7 @@ from executorch.backends.nxp.edge_passes.remove_as_strided_copy_nodes import ( RemoveUselessAsStridedCopyNodes, ) -from torch.fx.passes.infra.pass_manager import PassManager +from executorch.exir.pass_manager import PassManager class NeutronEdgePassManager(PassManager): diff --git a/backends/nxp/recipes/nxp_recipe_provider.py b/backends/nxp/recipes/nxp_recipe_provider.py new file mode 100644 index 00000000000..afef49d1cc7 --- /dev/null +++ b/backends/nxp/recipes/nxp_recipe_provider.py @@ -0,0 +1,353 @@ +# Copyright 2026 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 dataclasses import dataclass +from functools import partial +from typing import Any, Callable, cast, Iterable, Optional, Sequence + +import torch + +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec +from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass +from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import ( + NeutronEdgePassManager, +) +from executorch.backends.nxp.edge_passes.remove_additional_quantize_dequantize_nodes_pass import ( + RemoveAdditionalQDQClustersPass, +) +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 ( + core_aten_ops_exception_list, + generate_neutron_compile_spec, +) +from executorch.backends.nxp.quantizer.utils import calibrate_and_quantize +from executorch.backends.nxp.recipes.nxp_recipe_types import NXP_BACKEND, NXPRecipeType +from executorch.backends.nxp.tests.executorch_pipeline import ( + _get_default_quantizer, + get_random_calibration_inputs, + GetCalibrationInputsFn, + handle_kernel_selection, + ModelInputSpec, + to_model_input_spec, +) +from executorch.exir import EdgeCompileConfig, EdgeProgramManager, ExportedProgram + +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.backend.partitioner import Partitioner +from executorch.export import ( + BackendRecipeProvider, + ExportRecipe, + LoweringRecipe, + QuantizationRecipe, + RecipeType, +) +from torchao.quantization.pt2e.quantizer import Quantizer + + +class NeutronEdgePassManagerWrapper: + def __init__(self, passes: list[NeutronEdgePass] | None = None): + self.neutron_edge_pass_manager = NeutronEdgePassManager(passes) + + def __call__(self, s: str, epm: EdgeProgramManager) -> NeutronEdgePassManager: + return self.neutron_edge_pass_manager + + +NEUTRON_RECIPE_CONFIG_KEY = "neutron_recipe_config" + + +@dataclass +class NeutronRecipeConfig: + """Configuration shared by all NXP recipe types. + + Parameters that vary the *type* of export (PTQ vs QAT, delegate vs no-delegate) + are expressed by choosing a different NXPRecipeType rather than by flags here. + + Attributes: + input_spec: Model input description. Accepts a single shape tuple, a list of + shape tuples (one per input), or a list of ModelInputSpec objects. + target: Neutron hardware target string. Default: "imxrt700". + operators_not_to_delegate: Optional list of op names excluded from NPU delegation. + intermediates_dir: Optional directory to dump intermediate compilation artefacts. + get_quantizer_fn: Optional factory that returns a custom Quantizer. When None, + the default NeutronQuantizer is used. + get_calibration_inputs_fn: Optional function that, given the input_spec, returns + calibration input samples. When None, random inputs are + used. + train_fn: QAT training callback. Required when using INT8_QAT_NEUTRON. + custom_delegation_options: Optional fine-grained control over which ops are + delegated. Default: CustomDelegationOptions(). + remove_quant_io_ops: If True, remove quantize/dequantize ops at the IO boundary + (useful for integer-IO deployments). + use_quant_state_dict: If False, the post-quantization parameter values are not + passed to NeutronPartitioner. + use_neutron_for_format_conversion: Whether Neutron handles data-format conversion. + fetch_constants_to_sram: Place constant tensors in SRAM on the target. + dump_kernel_selection_code: Emit kernel-selection files after compilation. + use_profiling: Enable execution profiling / ETRecord generation. + """ + + input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]] + target: str = "imxrt700" + operators_not_to_delegate: list[str] = None + intermediates_dir: str | None = None + get_quantizer_fn: Callable[[], Quantizer] | None = None + get_calibration_inputs_fn: GetCalibrationInputsFn | None = None + train_fn: Callable[[torch.fx.GraphModule], None] | None = None + custom_delegation_options: CustomDelegationOptions | None = None + remove_quant_io_ops: bool = False + use_quant_state_dict: bool = True + use_neutron_for_format_conversion: bool = True + fetch_constants_to_sram: bool = False + dump_kernel_selection_code: bool = False + use_profiling: bool = False + + +class NXPRecipeProvider(BackendRecipeProvider): + + @property + def backend_name(self) -> str: + return NXP_BACKEND + + def get_supported_recipes(self) -> Sequence[RecipeType]: + return list(NXPRecipeType) + + def create_recipe( + self, recipe_type: RecipeType, **kwargs: Any + ) -> Optional[ExportRecipe]: + if recipe_type not in self.get_supported_recipes(): + logging.warning(f"NXP backend: Recipe `{recipe_type}` is not valid.") + return None + + rc = cast(NeutronRecipeConfig, kwargs.get(NEUTRON_RECIPE_CONFIG_KEY)) + if rc is None: + raise KeyError( + f"NXP backend: create_recipe() requires `{NEUTRON_RECIPE_CONFIG_KEY}=`." + ) + + if rc.custom_delegation_options is None: + rc.custom_delegation_options = CustomDelegationOptions() + + rc.input_spec = to_model_input_spec(rc.input_spec) + + match recipe_type: + case NXPRecipeType.INT8_PTQ_NEUTRON: + return self._build_recipe(recipe_type, rc, is_qat=False, delegate=True) + case NXPRecipeType.INT8_QAT_NEUTRON: + if rc.train_fn is None: + raise ValueError( + "NXP backend: INT8_QAT_NEUTRON requires train_fn in NeutronRecipeConfig." + ) + return self._build_recipe(recipe_type, rc, is_qat=True, delegate=True) + case NXPRecipeType.INT8_PTQ_NO_DELEGATE: + return self._build_recipe(recipe_type, rc, is_qat=False, delegate=False) + case _: + raise NotImplementedError( + f"NXP backend: Recipe `{recipe_type}` is not supported." + ) + + def _build_recipe( + self, + recipe_type: NXPRecipeType, + rc: NeutronRecipeConfig, + *, + is_qat: bool, + delegate: bool, + ) -> ExportRecipe: + neutron_target_spec = NeutronTargetSpec(rc.target) + + if rc.get_quantizer_fn is None: + rc.get_quantizer_fn = partial( + _get_default_quantizer, neutron_target_spec, is_qat + ) + + quantization_recipe = _build_quantization_recipe(rc, is_qat) + compile_spec = generate_neutron_compile_spec( + rc.target, + intermediates_dir=rc.intermediates_dir, + operators_not_to_delegate=rc.operators_not_to_delegate, + use_neutron_for_format_conversion=rc.use_neutron_for_format_conversion, + fetch_constants_to_sram=rc.fetch_constants_to_sram, + dump_kernel_selection_code=rc.dump_kernel_selection_code, + use_profiling=rc.use_profiling, + ) + lowering_recipe = _build_lowering_recipe( + compile_spec, neutron_target_spec, rc, delegate=delegate + ) + + return ExportRecipe( + name=recipe_type.value, + quantization_recipe=quantization_recipe, + lowering_recipe=lowering_recipe, + ) + + +# --------------------------------------------------------------------------- +# Module-level builder helpers +# --------------------------------------------------------------------------- + + +def _build_quantization_recipe( + rc: NeutronRecipeConfig, is_qat: bool +) -> QuantizationRecipe: + """Build the QuantizationRecipe for PTQ or QAT. + + PTQ uses the standard QuantizeStage flow (prepare_pt2e -> calibrate -> convert_pt2e) + driven by calibration_inputs_fn. quantize_fn is intentionally left None so that + multiple PTQ recipes can be combined with ExportRecipe.combine(). + + QAT requires a non-standard sequence (BN-fusion passes interleaved with training), + so it delegates to calibrate_and_quantize() via quantize_fn. A recipe with + quantize_fn set cannot be combined with other recipes. + """ + _input_spec = rc.input_spec + _calibration_fn = rc.get_calibration_inputs_fn or get_random_calibration_inputs + _quantizer = rc.get_quantizer_fn() + + if is_qat: + _train_fn = rc.train_fn + + def _quantize_fn( # noqa: E306 + graph_module: torch.fx.GraphModule, + ) -> torch.fx.GraphModule: + return calibrate_and_quantize( + model=graph_module, + calibration_inputs=_calibration_fn(_input_spec), + quantizer=_quantizer, + is_qat=True, + train_fn=_train_fn, + ) + + return QuantizationRecipe( + quantizers=[_quantizer], + quantize_fn=_quantize_fn, + ) + else: + + def _calibration_inputs_fn() -> Iterable: + return _calibration_fn(_input_spec) + + return QuantizationRecipe( + quantizers=[_quantizer], + calibration_inputs_fn=_calibration_inputs_fn, + ) + + +def _build_lowering_recipe( + compile_spec: list[CompileSpec], + neutron_target_spec: NeutronTargetSpec, + rc: NeutronRecipeConfig, + *, + delegate: bool, +) -> LoweringRecipe: + """Build the LoweringRecipe, optionally including NPU delegation.""" + partitioners = _build_partitioners(compile_spec, neutron_target_spec, rc, delegate) + pre_partitioning_callback = _build_pre_partitioning_callback(rc) + post_partitioning_transforms = _build_post_partitioning_transforms(rc) + + # The edge pass manager must be wrapped: EdgeTransformAndLowerStage calls + # edge_transform_passes with (method_name, ep) and expects a PassManager back. + return LoweringRecipe( + partitioners=partitioners, + edge_transform_passes=[NeutronEdgePassManagerWrapper()], + edge_compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _core_aten_ops_exception_list=core_aten_ops_exception_list, + ), + pre_partitioning_callback=pre_partitioning_callback, + post_partitioning_transforms=post_partitioning_transforms, + ) + + +def _build_partitioners( + compile_spec: list[CompileSpec], + neutron_target_spec: NeutronTargetSpec, + rc: NeutronRecipeConfig, + delegate: bool, +) -> list: + """Create the NeutronPartitioner list. Empty when delegate=False.""" + if not delegate: + return [] + return [ + NeutronPartitioner( + compile_spec, + neutron_target_spec, + rc.custom_delegation_options, + preserve_ops=[torch.ops.aten.prelu.default], + ) + ] + + +def _build_pre_partitioning_callback(rc: NeutronRecipeConfig): + """Return a callback that assigns the post-quantization state_dict to NeutronPartitioner. + + NeutronPartitioner requires static parameter data. Since the partitioner is instantiated + during recipe creation (before model data is available), assignment is deferred to a + callback invoked just before partitioning. + """ + _use_quant_state_dict = rc.use_quant_state_dict + + def _callback( + _partitioners: list[Partitioner] | None, + programs: dict[str, ExportedProgram], + ) -> None: + if _partitioners is None: + return + + if _use_quant_state_dict: + post_quant_state_dict: dict | None = {} + for _, program in programs.items(): + post_quant_state_dict.update(program.state_dict) + else: + post_quant_state_dict = None + + for _partitioner in _partitioners: + if isinstance(_partitioner, NeutronPartitioner): + _partitioner.post_quantization_state_dict = post_quant_state_dict + + return _callback + + +def _build_post_partitioning_transforms(rc: NeutronRecipeConfig) -> list: + """Build the list of post-partitioning EdgeProgramManager transforms. + + These mirror what the imperative pipeline did after to_edge_transform_and_lower: + - RemoveIOQuantOpsPass (optional, when remove_quant_io_ops=True) + - RemoveAdditionalQDQClustersPass (always applied) + - handle_kernel_selection side-effect (optional, when dump_kernel_selection_code=True) + """ + transforms = [] + + if rc.remove_quant_io_ops: + + def _remove_io_quant_ops(epm: EdgeProgramManager) -> EdgeProgramManager: + return epm.transform([RemoveIOQuantOpsPass(edge_program_manager=epm)]) + + transforms.append(_remove_io_quant_ops) + + def _remove_additional_qdq_clusters(epm: EdgeProgramManager) -> EdgeProgramManager: + return epm.transform( + NeutronEdgePassManager([RemoveAdditionalQDQClustersPass()]) + ) + + transforms.append(_remove_additional_qdq_clusters) + + if rc.dump_kernel_selection_code: + + def _handle_kernel_selection_transform( + epm: EdgeProgramManager, + ) -> EdgeProgramManager: + handle_kernel_selection() + return epm + + transforms.append(_handle_kernel_selection_transform) + + return transforms diff --git a/backends/nxp/recipes/nxp_recipe_types.py b/backends/nxp/recipes/nxp_recipe_types.py new file mode 100644 index 00000000000..1bef45922fe --- /dev/null +++ b/backends/nxp/recipes/nxp_recipe_types.py @@ -0,0 +1,35 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.export import RecipeType + + +NXP_BACKEND: str = "nxp" + + +class NXPRecipeType(RecipeType): + """NXP-specific recipe types for Neutron NPU export. + + Choose the recipe that matches your intended export configuration: + - INT8_PTQ_NEUTRON: standard post-training quantization, delegates to Neutron NPU. + - INT8_QAT_NEUTRON: quantization-aware training flow, delegates to Neutron NPU. + - INT8_PTQ_NO_DELEGATE: PTQ without NPU delegation (useful for debugging or CPU-only deployment). + """ + + # INT8 static PTQ (weights + activations). Calibration dataset required. + # Applicable operators are delegated to the Neutron NPU. + INT8_PTQ_NEUTRON = "nxp_int8_ptq_neutron" + + # INT8 QAT flow. Requires train_fn in NeutronRecipeConfig. + # Applicable operators are delegated to the Neutron NPU. + INT8_QAT_NEUTRON = "nxp_int8_qat_neutron" + + # INT8 PTQ without NPU delegation. Produces a quantized graph that runs on CPU. + # Useful for accuracy evaluation or debugging before enabling delegation. + INT8_PTQ_NO_DELEGATE = "nxp_int8_ptq_no_delegate" + + @classmethod + def get_backend_name(cls) -> str: + return NXP_BACKEND diff --git a/backends/nxp/tests/executorch_pipeline.py b/backends/nxp/tests/executorch_pipeline.py index 1309e019428..5f1e7c64b10 100644 --- a/backends/nxp/tests/executorch_pipeline.py +++ b/backends/nxp/tests/executorch_pipeline.py @@ -47,6 +47,7 @@ from torch import memory_format, nn from torch.export import export from torchao.quantization.pt2e.quantizer import Quantizer +from typing_extensions import deprecated neutron_target_spec = NeutronTargetSpec(target="imxrt700") @@ -174,6 +175,10 @@ def _get_example_input( return tuple(example_input) +@deprecated( + "NXP backend: The imperative lowering functions will be removed in the future. " + "Please use the recipe-based approach instead." +) def to_quantized_edge_program( model: torch.nn.Module, input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]], @@ -268,6 +273,10 @@ def to_quantized_edge_program( return edge_program_manager +@deprecated( + "NXP backend: The imperative lowering functions will be removed in the future. " + "Please use the recipe-based approach instead." +) def to_quantized_executorch_program( model: torch.nn.Module, input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]], diff --git a/backends/nxp/tests/generic_tests/test_recipe_export.py b/backends/nxp/tests/generic_tests/test_recipe_export.py new file mode 100644 index 00000000000..89dd1aad575 --- /dev/null +++ b/backends/nxp/tests/generic_tests/test_recipe_export.py @@ -0,0 +1,387 @@ +# Copyright 2026 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 pytest +import torch +import torch.nn + +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) +from executorch.backends.nxp.recipes.nxp_recipe_provider import ( + NEUTRON_RECIPE_CONFIG_KEY, + NeutronRecipeConfig, + NXPRecipeProvider, +) +from executorch.backends.nxp.recipes.nxp_recipe_types import NXPRecipeType +from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec +from executorch.backends.nxp.tests.executors import ( + graph_contains_any, + graph_contains_any_of_ops, +) +from executorch.backends.nxp.tests.ops_aliases import ExecutorchDelegateCall +from executorch.export import export +from executorch.export.recipe import ExportRecipe +from torch._inductor.lowering import quantized_decomposed + + +class SimpleCNN(torch.nn.Module): + def __init__(self, channels=3): + super().__init__() + self.conv = torch.nn.Conv2d(channels, channels, kernel_size=3) + + def forward(self, x): + x = self.conv(x) + x = torch.relu(x) + x = x.reshape(1, -1) + x = x + x + return x + + +class ConvBnRelu(torch.nn.Module): + """Conv + BatchNorm + ReLU - exercises QAT BN-fusion paths.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 3, kernel_size=3) + self.bn = torch.nn.BatchNorm2d(3) + + def forward(self, x): + return torch.relu(self.bn(self.conv(x))) + + +INPUT_SHAPE = (1, 3, 8, 8) + + +def _run_export( + model, rc, recipe_type=NXPRecipeType.INT8_PTQ_NEUTRON, input_shape=INPUT_SHAPE +): + example_inputs = [(torch.randn(input_shape),)] + recipe = NXPRecipeProvider().create_recipe(recipe_type, neutron_recipe_config=rc) + return export(model, example_inputs=example_inputs, export_recipe=recipe) + + +def _get_graph(sess): + return sess.get_edge_program_manager().exported_program().graph + + +def test_ptq_neutron_basic(): + """Baseline PTQ: whole model delegated, IO is quantized.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE) + sess = _run_export(model, rc) + graph = _get_graph(sess) + + assert graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + def is_cnn_op(n): + return any(op in n.name.lower() for op in ["conv", "relu", "view", "add"]) + + assert not graph_contains_any(graph, is_cnn_op) + + nodes = list(graph.nodes) + assert nodes[2].target == quantized_decomposed.quantize_per_tensor.out + assert nodes[-2].target == quantized_decomposed.dequantize_per_tensor.out + + +class TestInt8PTQNoDelegate: + + def test__basic(self): + """INT8_PTQ_NO_DELEGATE: model is quantized but no delegate call appears.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE) + sess = _run_export(model, rc, recipe_type=NXPRecipeType.INT8_PTQ_NO_DELEGATE) + graph = _get_graph(sess) + + assert not graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + def is_cnn_op(n): + return any(op in n.name.lower() for op in ["conv", "relu", "view", "add"]) + + # With no delegation, original ops should be visible in the graph. + assert graph_contains_any(graph, is_cnn_op) + + def test__recipe_has_empty_partitioners(self): + """INT8_PTQ_NO_DELEGATE recipe has an empty partitioner list.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NO_DELEGATE, neutron_recipe_config=rc + ) + assert recipe.lowering_recipe.partitioners == [] + + +class TestInt8QAT: + def test__requires_train_fn(self): + """INT8_QAT_NEUTRON raises ValueError when train_fn is missing.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) # no train_fn + with pytest.raises(ValueError, match="train_fn"): + NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_QAT_NEUTRON, neutron_recipe_config=rc + ) + + def test__recipe_structure(self): + """INT8_QAT_NEUTRON recipe has quantize_fn populated (QAT uses non-standard sequence).""" + + def dummy_train(m): + pass + + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=dummy_train) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_QAT_NEUTRON, neutron_recipe_config=rc + ) + qr = recipe.quantization_recipe + assert qr.quantize_fn is not None + assert ( + qr.calibration_inputs_fn is None + ) # QAT uses quantize_fn, not calibration_inputs_fn + + +class TestNeutronRecipeConfigFlags: + def test_operators_not_to_delegate(self): + """Ops listed in operators_not_to_delegate are not lowered to Neutron.""" + model = SimpleCNN() + rc = NeutronRecipeConfig( + INPUT_SHAPE, operators_not_to_delegate=["aten::convolution"] + ) + sess = _run_export(model, rc) + graph = _get_graph(sess) + + assert graph_contains_any_of_ops( + graph, [torch.ops.aten.convolution.out] + ) # Convolution was not delegated. + assert graph_contains_any_of_ops( + graph, [ExecutorchDelegateCall] + ) # Other operators were delegated. + + def _is_relu_add_or_view(n: torch.fx.Node) -> bool: + return any(op in n.name.lower() for op in ["relu", "add", "view"]) + + assert not graph_contains_any(graph, _is_relu_add_or_view) + + def test_remove_quant_io_ops(self): + """remove_quant_io_ops=True: no quantize op at the IO boundary.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, remove_quant_io_ops=True) + sess = _run_export(model, rc) + graph = _get_graph(sess) + nodes = list(graph.nodes) + + real_nodes = [n for n in nodes if n.op not in ("placeholder", "output")] + assert real_nodes[0].target != quantized_decomposed.quantize_per_tensor.out + assert real_nodes[-1].target != quantized_decomposed.dequantize_per_tensor.out + assert real_nodes[-1].meta["val"].dtype == torch.int8 + placeholder_nodes = [n for n in nodes if n.op == "placeholder"] + assert placeholder_nodes[0].name == "x" # Main input + assert placeholder_nodes[0].meta["val"].dtype == torch.int8 + + def test_use_quant_state_dict_false(self): + """use_quant_state_dict=False: NeutronPartitioner gets None state_dict, no crash.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, use_quant_state_dict=False) + sess = _run_export(model, rc) + assert sess.get_edge_program_manager() is not None + + def test_custom_delegation_options_explicit(self): + """Explicitly provided CustomDelegationOptions is forwarded correctly.""" + model = SimpleCNN() + opts = CustomDelegationOptions() + rc = NeutronRecipeConfig(INPUT_SHAPE, custom_delegation_options=opts) + sess = _run_export(model, rc) + assert sess.get_edge_program_manager() is not None + + def test_custom_quantizer_fn(self): + """get_quantizer_fn overrides the default NeutronQuantizer.""" + from executorch.backends.nxp.backend.neutron_target_spec import ( + NeutronTargetSpec, + ) + from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer + + custom_quantizer_called = [] + + def my_quantizer_fn(): + q = NeutronQuantizer(NeutronTargetSpec("imxrt700")) + custom_quantizer_called.append(True) + return q + + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, get_quantizer_fn=my_quantizer_fn) + sess = _run_export(model, rc) + assert custom_quantizer_called, "Custom quantizer factory was not called." + assert sess.get_edge_program_manager() is not None + + def test_custom_calibration_inputs_fn(self): + """get_calibration_inputs_fn is called for calibration.""" + calibration_called = [] + + def my_calibration_fn(input_spec): + calibration_called.append(True) + return [ + tuple(torch.randn(s.shape, dtype=s.dtype) for s in input_spec) + for _ in range(2) + ] + + model = SimpleCNN() + rc = NeutronRecipeConfig( + INPUT_SHAPE, get_calibration_inputs_fn=my_calibration_fn + ) + sess = _run_export(model, rc) + assert calibration_called, "Custom calibration function was not called." + assert sess.get_edge_program_manager() is not None + + def test_use_neutron_for_format_conversion_false(self): + """use_neutron_for_format_conversion=False still produces a valid export.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, use_neutron_for_format_conversion=False) + sess = _run_export(model, rc) + assert sess.get_edge_program_manager() is not None + + def test_target_explicit(self): + """Specifying fake target to make sure an error is raised.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, target="FAKE") + with pytest.raises(ValueError, match="`FAKE` is not a valid target"): + _run_export(model, rc) + + +class TestInputSpecForms: + def test__single_tuple(self): + """input_spec as a plain shape tuple works.""" + model = SimpleCNN() + sess = _run_export(model, NeutronRecipeConfig((1, 3, 8, 8))) + assert sess.get_edge_program_manager() is not None + + def test__list_of_tuples(self): + """input_spec as list of shape tuples works.""" + model = SimpleCNN() + sess = _run_export(model, NeutronRecipeConfig([(1, 3, 8, 8)])) + assert sess.get_edge_program_manager() is not None + + def test__model_input_spec(self): + """input_spec as list of ModelInputSpec objects works.""" + model = SimpleCNN() + sess = _run_export(model, NeutronRecipeConfig([ModelInputSpec((1, 3, 8, 8))])) + assert sess.get_edge_program_manager() is not None + + +class TestErrorHandling: + def test_create_recipe_missing_config_key(self): + """create_recipe without neutron_recipe_config kwarg raises KeyError.""" + with pytest.raises(KeyError, match=NEUTRON_RECIPE_CONFIG_KEY): + NXPRecipeProvider().create_recipe(NXPRecipeType.INT8_PTQ_NEUTRON) + + def test_create_recipe_invalid_recipe_type(self): + """create_recipe with an unsupported recipe type returns None with a warning.""" + from executorch.export.recipe import RecipeType + + class FakeRecipeType(RecipeType): + FAKE = "fake" + + @classmethod + def get_backend_name(cls): + return "fake_backend" + + rc = NeutronRecipeConfig(INPUT_SHAPE) + result = NXPRecipeProvider().create_recipe( + FakeRecipeType.FAKE, neutron_recipe_config=rc + ) + assert result is None + + +class TestRecipeStructureValidation: + + def test_ptq_neutron_recipe_structure(self): + """INT8_PTQ_NEUTRON recipe: correct quantizer, partitioner, no quantize_fn.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + assert recipe.quantization_recipe is not None + assert len(recipe.quantization_recipe.quantizers) == 1 + # PTQ must NOT use quantize_fn so that multiple PTQ recipes can be combined. + assert recipe.quantization_recipe.quantize_fn is None + assert recipe.quantization_recipe.calibration_inputs_fn is not None + assert recipe.lowering_recipe.partitioners is not None + assert len(recipe.lowering_recipe.partitioners) == 1 + + def test_ptq_neutron_recipe_name(self): + """INT8_PTQ_NEUTRON recipe has the expected name.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + assert recipe.name == NXPRecipeType.INT8_PTQ_NEUTRON.value + + +class TestRecipeCombination: + def test__chains_pre_partitioning_callbacks(self): + """Combining two NXP recipes chains both pre_partitioning_callbacks.""" + recipe1 = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, + neutron_recipe_config=NeutronRecipeConfig(INPUT_SHAPE), + ) + recipe2 = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, + neutron_recipe_config=NeutronRecipeConfig(INPUT_SHAPE), + ) + combined = ExportRecipe.combine([recipe1, recipe2]) + assert combined.lowering_recipe.pre_partitioning_callback is not None + # Calling the combined callback should not raise. + combined.lowering_recipe.pre_partitioning_callback(None, {}) + + def test__combine_raises_for_qat_recipe(self): + """ExportRecipe.combine() raises ValueError when any recipe has quantize_fn set.""" + + def dummy_train(m): + pass + + qat_recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_QAT_NEUTRON, + neutron_recipe_config=NeutronRecipeConfig( + INPUT_SHAPE, train_fn=dummy_train + ), + ) + ptq_recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, + neutron_recipe_config=NeutronRecipeConfig(INPUT_SHAPE), + ) + with pytest.raises(ValueError, match="quantize_fn"): + ExportRecipe.combine([ptq_recipe, qat_recipe]) + + def test__collects_post_partitioning_transforms(self): + """Combining two NXP recipes collects post_partitioning_transforms from both.""" + recipe1 = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, + neutron_recipe_config=NeutronRecipeConfig(INPUT_SHAPE), + ) + recipe2 = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, + neutron_recipe_config=NeutronRecipeConfig(INPUT_SHAPE), + ) + n1 = len(recipe1.lowering_recipe.post_partitioning_transforms) + n2 = len(recipe2.lowering_recipe.post_partitioning_transforms) + combined = ExportRecipe.combine([recipe1, recipe2]) + assert len(combined.lowering_recipe.post_partitioning_transforms) == n1 + n2 + + +class TestPostPartitioningTransforms: + def test__executed(self): + """post_partitioning_transforms are called after partitioning.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + + transform_called = [] + + def tracking_transform(epm): + transform_called.append(True) + return epm + + recipe.lowering_recipe.post_partitioning_transforms = [tracking_transform] + + example_inputs = [(torch.randn(INPUT_SHAPE),)] + export(model, example_inputs=example_inputs, export_recipe=recipe) + assert transform_called, "post_partitioning_transforms were not executed." diff --git a/export/recipe.py b/export/recipe.py index c1227384de4..34f46304597 100644 --- a/export/recipe.py +++ b/export/recipe.py @@ -1,14 +1,16 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # Copyright 2025 Arm Limited and/or its affiliates. +# Copyright 2026 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 abc import ABCMeta, abstractmethod from dataclasses import dataclass from enum import Enum, EnumMeta -from typing import Callable, List, Optional +from typing import Callable, Iterable, List, Optional import torch from executorch.exir import EdgeProgramManager, ExportedProgram @@ -94,10 +96,22 @@ class QuantizationRecipe: quantizers: Optional list of quantizers for model quantization ao_quantization_configs: Optional list of AOQuantizationConfig objects that pair AOBaseConfig with optional filter functions + calibration_inputs_fn: Optional zero-argument callable that returns an iterable of + calibration input tuples. When None the example_inputs passed + to the export session are used directly. + Ignored when quantize_fn is set. + quantize_fn: Optional callable that replaces the default prepare/calibrate/convert flow + in QuantizeStage. Use this when the built-in flow does not support the + required quantization sequence. Recipes that set this field cannot be + combined with other recipes via ExportRecipe.combine(). """ quantizers: Optional[List[Quantizer]] = None ao_quantization_configs: Optional[List[AOQuantizationConfig]] = None + calibration_inputs_fn: Optional[ + Callable[[], "Iterable[tuple[torch.Tensor, ...]]"] + ] = None + quantize_fn: Optional[Callable[[torch.fx.GraphModule], torch.fx.GraphModule]] = None def get_quantizers(self) -> Optional[List[Quantizer]]: """ @@ -124,6 +138,13 @@ class LoweringRecipe: edge_manager_transform_passes: Optional list of callables that take EdgeProgramManager as argument and return passes to be applied. Applied sequentially after TO_EDGE stage. edge_compile_config: Optional edge compilation configuration + pre_partitioning_callback: Optional callable invoked just before partitioning with + `(partitioners, programs)` arguments. + post_partitioning_transforms: Optional list of callables each with signature + `(EdgeProgramManager) -> EdgeProgramManager`, applied in + order after partitioning. Use this for graph transforms that + require the fully partitioned graph, or for side-effect + operations that must run after delegation. """ partitioners: Optional[List[Partitioner]] = None @@ -136,6 +157,12 @@ class LoweringRecipe: ) = None # pyre-ignore[11]: Type not defined edge_compile_config: Optional[EdgeCompileConfig] = None + pre_partitioning_callback: Optional[ + Callable[[Optional[list[Partitioner]], dict[str, ExportedProgram]], None] + ] = None + post_partitioning_transforms: Optional[ + List[Callable[["EdgeProgramManager"], "EdgeProgramManager"]] + ] = None @experimental( @@ -245,12 +272,25 @@ def _combine_recipes( # noqa: C901 Returns: Combined ExportRecipe for multi-backend deployment """ + for recipe in backend_recipes: + if ( + recipe.quantization_recipe is not None + and recipe.quantization_recipe.quantize_fn is not None + ): + raise ValueError( + f"Recipe '{recipe.name}' uses a custom quantize_fn and cannot be " + "combined with other recipes. quantize_fn calls convert_pt2e internally " + "and the resulting quantized graph cannot be merged with other recipes." + ) + # Extract components from individual recipes all_partitioners = [] all_quantizers = [] all_ao_quantization_configs = [] all_pre_edge_passes = [] all_transform_passes = [] + all_pre_partitioning_callbacks = [] + all_post_partitioning_passes = [] combined_backend_config = None for recipe in backend_recipes: @@ -268,6 +308,24 @@ def _combine_recipes( # noqa: C901 recipe.lowering_recipe.edge_transform_passes ) + # Collect pre-partitioning callbacks - all will be chained + if ( + recipe.lowering_recipe + and recipe.lowering_recipe.pre_partitioning_callback + ): + all_pre_partitioning_callbacks.append( + recipe.lowering_recipe.pre_partitioning_callback + ) + + # Collect post-partitioning transforms + if ( + recipe.lowering_recipe + and recipe.lowering_recipe.post_partitioning_transforms + ): + all_post_partitioning_passes.extend( + recipe.lowering_recipe.post_partitioning_transforms + ) + # Collect for quantize stage if quantization_recipe := recipe.quantization_recipe: # Collect PT2E quantizers @@ -296,9 +354,25 @@ def _combine_recipes( # noqa: C901 ), ) + # Chain all pre-partitioning callbacks into a single one that calls each in order + combined_pre_partitioning_callback = None + if all_pre_partitioning_callbacks: + _cbs = all_pre_partitioning_callbacks + + def _chained_pre_partitioning_callback(partitioners, programs): + for cb in _cbs: + cb(partitioners, programs) + + combined_pre_partitioning_callback = _chained_pre_partitioning_callback + # Create combined lowering recipe combined_lowering_recipe = None - if all_partitioners or all_transform_passes: + if ( + all_partitioners + or all_transform_passes + or combined_pre_partitioning_callback + or all_post_partitioning_passes + ): edge_compile_config = None for recipe in backend_recipes: if ( @@ -314,6 +388,8 @@ def _combine_recipes( # noqa: C901 all_transform_passes if all_transform_passes else None ), edge_compile_config=edge_compile_config or EdgeCompileConfig(), + pre_partitioning_callback=combined_pre_partitioning_callback, + post_partitioning_transforms=all_post_partitioning_passes or None, ) recipe_name = recipe_name or "_".join( diff --git a/export/stages.py b/export/stages.py index dbc4703faaa..562db0432b3 100644 --- a/export/stages.py +++ b/export/stages.py @@ -1,6 +1,7 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # Copyright 2025 Arm Limited and/or its affiliates. +# Copyright 2026 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -240,6 +241,13 @@ def run(self, artifact: PipelineArtifact) -> None: # Use PassManager directly if found, otherwise use dict final_passes = pass_manager if pass_manager else transform_passes + if ( + hasattr(er := artifact.context["export_recipe"], "lowering_recipe") + and er.lowering_recipe is not None + and (callback := er.lowering_recipe.pre_partitioning_callback) is not None + ): + callback(self._partitioners, artifact.data) + with validation_disabled(): edge_program_manager = to_edge_transform_and_lower( exported_programs, @@ -250,6 +258,16 @@ def run(self, artifact: PipelineArtifact) -> None: generate_etrecord=generate_etrecord, ) + # Apply post-partitioning transforms if specified in the lowering recipe. These are used for graph transforms + # that require the fully partitioned graph or for side effect operations. + if ( + hasattr(er := artifact.context.get("export_recipe"), "lowering_recipe") + and er.lowering_recipe is not None + and (transforms := er.lowering_recipe.post_partitioning_transforms) + ): + for transform in transforms: + edge_program_manager = transform(edge_program_manager) + delegation_info = get_delegation_info( edge_program_manager.exported_program().graph_module ) @@ -411,9 +429,12 @@ def _get_quantizer_for_prepare_pt2e(self, quantizers: List[Any]): raise ValueError("No quantizers detected") def run(self, artifact: PipelineArtifact) -> None: - if not self._quantization_recipe or not self._quantization_recipe.quantizers: + if not self._quantization_recipe or ( + not self._quantization_recipe.quantizers + and self._quantization_recipe.quantize_fn is None + ): logging.info( - "Quantization recipe is invalid to run QunatizeStage, returning original model" + "Quantization recipe is invalid to run QuantizeStage, returning original model" ) self._artifact = artifact return @@ -431,18 +452,33 @@ def run(self, artifact: PipelineArtifact) -> None: f"Example inputs for method {method_name} not found or empty." ) + if isinstance(model, torch.nn.Module): + model.eval() + inputs = example_inputs[method_name][0] captured_graph = torch.export.export(model, inputs, strict=True).module() - quantizer = self._get_quantizer_for_prepare_pt2e( - self._quantization_recipe.quantizers # pyre-ignore - ) - prepared_model = prepare_pt2e(captured_graph, quantizer) - - for calibration_input in example_inputs[method_name]: - prepared_model(*calibration_input) + if (quantize_fn := self._quantization_recipe.quantize_fn) is not None: + # Use the backend-supplied quantization function directly. This allows backends to use non-standard + # quantization sequences (e.g. QAT with ordered BN-fusion passes around training and calibration) + # without exposing multiple ordered hook points in QuantizationRecipe. + quantized_model = quantize_fn(captured_graph) + else: + quantizer = self._get_quantizer_for_prepare_pt2e( + self._quantization_recipe.quantizers # pyre-ignore + ) + # PTQ flow: prepare_pt2e -> calibrate -> convert_pt2e. + # Use calibration_inputs_fn when provided, or fall back to the example_inputs. + m = prepare_pt2e(captured_graph, quantizer) + calibration_inputs = ( + self._quantization_recipe.calibration_inputs_fn() + if self._quantization_recipe.calibration_inputs_fn is not None + else example_inputs[method_name] + ) + for calibration_input in calibration_inputs: + m(*calibration_input) + quantized_model = convert_pt2e(m) - quantized_model = convert_pt2e(prepared_model) quantized_models[method_name] = quantized_model self._artifact = artifact.copy_with_new_data(quantized_models)