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..4aeb14a529e --- /dev/null +++ b/backends/nxp/recipes/nxp_recipe_provider.py @@ -0,0 +1,374 @@ +# 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.aten_passes.fuse_batch_norm_with_linear_pass import ( + FuseBatchNormWithLinearPass, +) +from executorch.backends.nxp.aten_passes.simulated_linear_bn_fusion_passes import ( + AddSimulatedLinearBatchNormFusionQATPass, + RemoveSimulatedLinearBatchNormFusionQATPass, +) +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.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.backends.transforms.quantize_fused_convbn_bias_pass import ( + QuantizeFusedConvBnBiasAtenPass, +) +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 default flow (prepare_pt2e -> calibrate -> + convert_pt2e) with calibration_inputs_fn supplying the calibration samples. + + QAT uses the same QuantizeStage flow with is_qat=True (prepare_qat_pt2e) and + the NXP-specific BN-fusion passes distributed across the four pass-list hooks: + - post_prepare_passes: AddSimulatedLinearBatchNormFusionQATPass + - post_train_passes: RemoveSimulatedLinearBatchNormFusionQATPass, + FuseBatchNormWithLinearPass + - post_calibration_passes: RemoveSimulatedLinearBatchNormFusionQATPass, + FuseBatchNormWithLinearPass (applied again after + observer-only calibration when train_fn is None) + - post_convert_passes: QuantizeFusedConvBnBiasAtenPass + """ + _input_spec = rc.input_spec + _calibration_fn = rc.get_calibration_inputs_fn or get_random_calibration_inputs + _quantizer = rc.get_quantizer_fn() + + def _calibration_inputs_fn() -> Iterable: + return _calibration_fn(_input_spec) + + if not is_qat: + return QuantizationRecipe( + quantizers=[_quantizer], + calibration_inputs_fn=_calibration_inputs_fn, + ) + + # QAT pass hooks -- keep NXP-specific pass classes out of shared modules + def _post_prepare(m: torch.fx.GraphModule) -> torch.fx.GraphModule: + return AddSimulatedLinearBatchNormFusionQATPass()(m).graph_module + + def _remove_simulated_bn_and_fuse(m: torch.fx.GraphModule) -> torch.fx.GraphModule: + m = RemoveSimulatedLinearBatchNormFusionQATPass()(m).graph_module + return FuseBatchNormWithLinearPass()(m).graph_module + + def _post_convert(m: torch.fx.GraphModule) -> torch.fx.GraphModule: + return QuantizeFusedConvBnBiasAtenPass( + default_zero_bias=False, symmetric_quant=True + )(m).graph_module + + return QuantizationRecipe( + quantizers=[_quantizer], + calibration_inputs_fn=_calibration_inputs_fn, + is_qat=True, + train_fn=rc.train_fn, + post_prepare_passes=[_post_prepare], + # Applied after training (or after prepare when no train_fn). + post_train_passes=[_remove_simulated_bn_and_fuse], + # Applied after calibration (only reached when train_fn is None). + post_calibration_passes=[_remove_simulated_bn_and_fuse], + post_convert_passes=[_post_convert], + ) + + +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 a186535e53d..21362b477ab 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..57d1edc3fbc --- /dev/null +++ b/backends/nxp/tests/generic_tests/test_recipe_export.py @@ -0,0 +1,429 @@ +# 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: is_qat=True, train_fn forwarded, pass hooks populated.""" + + 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.is_qat is True + assert qr.train_fn is dummy_train + assert qr.calibration_inputs_fn is not None + assert qr.post_prepare_passes is not None + assert qr.post_train_passes is not None + assert qr.post_calibration_passes is not None + assert qr.post_convert_passes is not None + + +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 QAT pass hooks.""" + 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 + assert recipe.quantization_recipe.is_qat is False + assert recipe.quantization_recipe.calibration_inputs_fn is not None + # PTQ does not use QAT pass hooks. + assert recipe.quantization_recipe.post_prepare_passes is None + assert recipe.quantization_recipe.post_convert_passes is 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), + ) + # Share the same calibration_inputs_fn so the single-owner guard does not raise. + shared_fn = recipe1.quantization_recipe.calibration_inputs_fn + recipe2.quantization_recipe.calibration_inputs_fn = shared_fn + + 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_merges_qat_pass_lists(self): + """Combining two QAT recipes concatenates all four quantization pass lists. + + Two recipes are built sharing the same train_fn and calibration_inputs_fn so + that _combine_recipes does not raise the single-owner guards. The pass counts + from both recipes must be present in the combined recipe. + """ + + def dummy_train(m): + pass + + def shared_calibration_fn(): + return [(torch.randn(INPUT_SHAPE),)] + + rc1 = NeutronRecipeConfig(INPUT_SHAPE, train_fn=dummy_train) + rc2 = NeutronRecipeConfig(INPUT_SHAPE, train_fn=dummy_train) + qat_recipe1 = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_QAT_NEUTRON, neutron_recipe_config=rc1 + ) + qat_recipe2 = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_QAT_NEUTRON, neutron_recipe_config=rc2 + ) + + # Override calibration_inputs_fn to be the same object on both recipes + # so the single-owner guard does not raise. + qat_recipe1.quantization_recipe.calibration_inputs_fn = shared_calibration_fn + qat_recipe2.quantization_recipe.calibration_inputs_fn = shared_calibration_fn + + qr1 = qat_recipe1.quantization_recipe + qr2 = qat_recipe2.quantization_recipe + + combined = ExportRecipe.combine([qat_recipe1, qat_recipe2]) + cqr = combined.quantization_recipe + + assert cqr.is_qat is True + assert len(cqr.post_prepare_passes) == len(qr1.post_prepare_passes) + len( + qr2.post_prepare_passes + ) + assert len(cqr.post_train_passes) == len(qr1.post_train_passes) + len( + qr2.post_train_passes + ) + assert len(cqr.post_calibration_passes) == len( + qr1.post_calibration_passes + ) + len(qr2.post_calibration_passes) + assert len(cqr.post_convert_passes) == len(qr1.post_convert_passes) + len( + qr2.post_convert_passes + ) + + 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), + ) + # Share the same calibration_inputs_fn so the single-owner guard does not raise. + shared_fn = recipe1.quantization_recipe.calibration_inputs_fn + recipe2.quantization_recipe.calibration_inputs_fn = shared_fn + + 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..ec51cbf9ab1 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 @@ -89,15 +91,60 @@ class QuantizationRecipe: Configuration recipe for quantization. This class holds the configuration parameters for quantizing a model. + The quantization sequence executed by QuantizeStage is: + + prepare_qat_pt2e / prepare_pt2e + -> post_prepare_passes + -> [train_fn] + -> post_train_passes + -> [calibration loop] (skipped when train_fn is provided) + -> post_calibration_passes + -> convert_pt2e + -> post_convert_passes + + For standard PTQ, all pass lists can be left None. For QAT or other + non-standard sequences, backends supply the appropriate passes via the + four pass-list fields. Attributes: - quantizers: Optional list of quantizers for model quantization + quantizers: Optional list of quantizers for model quantization. ao_quantization_configs: Optional list of AOQuantizationConfig objects that pair - AOBaseConfig with optional filter functions + 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 as calibration inputs. + is_qat: When True, QuantizeStage uses prepare_qat_pt2e instead of prepare_pt2e. + train_fn: Optional training callback invoked between post_prepare_passes and + post_train_passes. When provided the calibration loop is skipped. + post_prepare_passes: Optional list of GraphModule transforms applied after + prepare_pt2e / prepare_qat_pt2e. + post_train_passes: Optional list of GraphModule transforms applied after train_fn + (or after prepare, when no train_fn is provided). + post_calibration_passes: Optional list of GraphModule transforms applied after + the calibration loop. + post_convert_passes: Optional list of GraphModule transforms applied after + convert_pt2e. """ quantizers: Optional[List[Quantizer]] = None ao_quantization_configs: Optional[List[AOQuantizationConfig]] = None + calibration_inputs_fn: Optional[ + Callable[[], "Iterable[tuple[torch.Tensor, ...]]"] + ] = None + is_qat: bool = False + train_fn: Optional[Callable[[torch.fx.GraphModule], None]] = None + post_prepare_passes: Optional[ + List[Callable[[torch.fx.GraphModule], torch.fx.GraphModule]] + ] = None + post_train_passes: Optional[ + List[Callable[[torch.fx.GraphModule], torch.fx.GraphModule]] + ] = None + post_calibration_passes: Optional[ + List[Callable[[torch.fx.GraphModule], torch.fx.GraphModule]] + ] = None + post_convert_passes: Optional[ + List[Callable[[torch.fx.GraphModule], torch.fx.GraphModule]] + ] = None def get_quantizers(self) -> Optional[List[Quantizer]]: """ @@ -124,6 +171,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 +190,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( @@ -251,8 +311,20 @@ def _combine_recipes( # noqa: C901 all_ao_quantization_configs = [] all_pre_edge_passes = [] all_transform_passes = [] + all_edge_manager_transform_passes = [] + all_pre_partitioning_callbacks = [] + all_post_partitioning_passes = [] combined_backend_config = None + # Quantization pass lists accumulated across all recipes + all_post_prepare_passes = [] + all_post_train_passes = [] + all_post_calibration_passes = [] + all_post_convert_passes = [] + combined_is_qat = False + combined_train_fn = None + combined_calibration_inputs_fn = None + for recipe in backend_recipes: # Collect pre-edge transform passes if recipe.aten_transform_passes: @@ -262,23 +334,77 @@ def _combine_recipes( # noqa: C901 if recipe.lowering_recipe and recipe.lowering_recipe.partitioners: all_partitioners.extend(recipe.lowering_recipe.partitioners) - # Collect transform passes from lowering recipes + # Collect edge transform passes from lowering recipes if recipe.lowering_recipe and recipe.lowering_recipe.edge_transform_passes: all_transform_passes.extend( recipe.lowering_recipe.edge_transform_passes ) - # Collect for quantize stage - if quantization_recipe := recipe.quantization_recipe: - # Collect PT2E quantizers - if quantization_recipe.quantizers: - all_quantizers.extend(quantization_recipe.quantizers) + # Collect edge manager transform passes from lowering recipes + if ( + recipe.lowering_recipe + and recipe.lowering_recipe.edge_manager_transform_passes + ): + all_edge_manager_transform_passes.extend( + recipe.lowering_recipe.edge_manager_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 source transform configs - if quantization_recipe.ao_quantization_configs: - all_ao_quantization_configs.extend( - quantization_recipe.ao_quantization_configs - ) + # 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 qr := recipe.quantization_recipe: + if qr.quantizers: + all_quantizers.extend(qr.quantizers) + if qr.ao_quantization_configs: + all_ao_quantization_configs.extend(qr.ao_quantization_configs) + if qr.is_qat: + combined_is_qat = True + if qr.train_fn is not None: + if ( + combined_train_fn is not None + and combined_train_fn is not qr.train_fn + ): + raise ValueError( + f"Recipe '{recipe.name}' sets train_fn but another recipe already " + "provides one. Only one train_fn is allowed when combining recipes." + ) + combined_train_fn = qr.train_fn + if qr.calibration_inputs_fn is not None: + if ( + combined_calibration_inputs_fn is not None + and combined_calibration_inputs_fn + is not qr.calibration_inputs_fn + ): + raise ValueError( + f"Recipe '{recipe.name}' sets calibration_inputs_fn but another " + "recipe already provides one. Only one calibration_inputs_fn is " + "allowed when combining recipes." + ) + combined_calibration_inputs_fn = qr.calibration_inputs_fn + if qr.post_prepare_passes: + all_post_prepare_passes.extend(qr.post_prepare_passes) + if qr.post_train_passes: + all_post_train_passes.extend(qr.post_train_passes) + if qr.post_calibration_passes: + all_post_calibration_passes.extend(qr.post_calibration_passes) + if qr.post_convert_passes: + all_post_convert_passes.extend(qr.post_convert_passes) # Use the first backend config as base if combined_backend_config is None and recipe.executorch_backend_config: @@ -288,17 +414,49 @@ def _combine_recipes( # noqa: C901 # Create combined quantization recipe combined_quantization_recipe = None - if all_quantizers or all_ao_quantization_configs: + if ( + all_quantizers + or all_ao_quantization_configs + or combined_is_qat + or all_post_prepare_passes + or all_post_train_passes + or all_post_calibration_passes + or all_post_convert_passes + ): combined_quantization_recipe = QuantizationRecipe( quantizers=all_quantizers if all_quantizers else None, ao_quantization_configs=( all_ao_quantization_configs if all_ao_quantization_configs else None ), + calibration_inputs_fn=combined_calibration_inputs_fn, + is_qat=combined_is_qat, + train_fn=combined_train_fn, + post_prepare_passes=all_post_prepare_passes or None, + post_train_passes=all_post_train_passes or None, + post_calibration_passes=all_post_calibration_passes or None, + post_convert_passes=all_post_convert_passes or None, ) + # 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 all_edge_manager_transform_passes + or combined_pre_partitioning_callback + or all_post_partitioning_passes + ): edge_compile_config = None for recipe in backend_recipes: if ( @@ -313,7 +471,14 @@ def _combine_recipes( # noqa: C901 edge_transform_passes=( all_transform_passes if all_transform_passes else None ), + edge_manager_transform_passes=( + all_edge_manager_transform_passes + if all_edge_manager_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..39b7a5b0861 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. @@ -22,7 +23,15 @@ from torch import nn from torch._export.pass_base import PassType from torchao.quantization import quantize_ -from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +from torchao.quantization.pt2e import ( + move_exported_model_to_eval, + move_exported_model_to_train, +) +from torchao.quantization.pt2e.quantize_pt2e import ( + convert_pt2e, + prepare_pt2e, + prepare_qat_pt2e, +) from torchao.quantization.pt2e.quantizer import ( ComposableQuantizer, Quantizer as TorchAOPT2EQuantizer, @@ -240,6 +249,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 +266,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 ) @@ -410,16 +436,27 @@ def _get_quantizer_for_prepare_pt2e(self, quantizers: List[Any]): else: raise ValueError("No quantizers detected") + @staticmethod + def _apply_passes( + m: torch.fx.GraphModule, + passes: Optional[List[Callable[[torch.fx.GraphModule], torch.fx.GraphModule]]], + ) -> torch.fx.GraphModule: + """Apply a list of GraphModule -> GraphModule transforms in order.""" + for p in passes or []: + m = p(m) + return m + def run(self, artifact: PipelineArtifact) -> None: if not self._quantization_recipe or not self._quantization_recipe.quantizers: 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 assert isinstance(artifact.data, dict) + qr = self._quantization_recipe models = artifact.data example_inputs = artifact.get_context("example_inputs") @@ -431,19 +468,55 @@ 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 + qr.quantizers # pyre-ignore ) - prepared_model = prepare_pt2e(captured_graph, quantizer) - for calibration_input in example_inputs[method_name]: - prepared_model(*calibration_input) + # Phase 1: prepare + if qr.is_qat: + m = prepare_qat_pt2e(captured_graph, quantizer) + else: + m = prepare_pt2e(captured_graph, quantizer) + + m = self._apply_passes(m, qr.post_prepare_passes) + + # Phase 2: optional training. + # In QAT mode the model is in training mode after prepare_qat_pt2e, so + # move_exported_model_to_eval is always called (with or without a train_fn) + # to ensure eval mode before the subsequent passes and calibration. + if qr.train_fn is not None: + m = move_exported_model_to_train(m) + qr.train_fn(m) + + if qr.is_qat: + m = move_exported_model_to_eval(m) + + m = self._apply_passes(m, qr.post_train_passes) + + # Phase 3: calibration (skipped when a train_fn was provided) + if qr.train_fn is None: + calibration_inputs = ( + qr.calibration_inputs_fn() + if qr.calibration_inputs_fn is not None + else example_inputs[method_name] + ) + for calibration_input in calibration_inputs: + m(*calibration_input) + + m = self._apply_passes(m, qr.post_calibration_passes) + + # Phase 4: convert + m = convert_pt2e(m) + + m = self._apply_passes(m, qr.post_convert_passes) - quantized_model = convert_pt2e(prepared_model) - quantized_models[method_name] = quantized_model + quantized_models[method_name] = m self._artifact = artifact.copy_with_new_data(quantized_models)