NXP backend: Test avg_pool2d with new Neutron flow.

backends/nxp/backend/custom_delegation_options.py

@@ -22,3 +22,5 @@ class CustomDelegationOptions:
     #  not create any NeutronGraph that can be called. This is done by the partitioner itself, and is not handled by
     #  the individual node converters.
     allow_no_op_partitions: bool = False
+
+    use_new_flow_neutron_c: bool = False

backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py

@@ -1,4 +1,4 @@
-# Copyright 2025 NXP
+# Copyright 2025-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.
@@ -21,6 +21,8 @@
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
     average_pool_2d_options,
 )
+
+from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
 from torch.fx import Node
 from torch.nn import Parameter
 
@@ -53,6 +55,27 @@ def _is_supported_in_IR(
 
         return True
 
+    @staticmethod
+    def _is_supported_on_target(
+        node: Node,
+        neutron_target_spec: NeutronTargetSpec,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
+    ) -> bool:
+        kernel = node.args[1]
+        stride = node.args[2]
+
+        if custom_delegation_options.use_new_flow_neutron_c:
+            # Requirements specified by the new Neutron flow documentation.
+
+            if any(k > 4096 for k in kernel):
+                return False
+
+            if any(s > 4096 for s in stride):
+                return False
+
+        return True
+
     # noinspection PyMethodMayBeStatic
     def _convert_2d_avg_pool(
         self, kernel_size, stride, padding, t_op: tflite_model.Operator
@@ -85,10 +108,19 @@ def _convert_2d_avg_pool(
 
         return ops.flatten()
 
-    # AvgPool2d Node format: (Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, bool ceil_mode=False
-    #                         bool count_include_pad=True, int? divisor_override=None)
     def convert(self, node: Node):
-        """Convert 'avg_pool2d' operator to TFLite 'AveragePool2D'."""
+        """Convert 'avg_pool2d' operator to TFLite 'AveragePool2D'.
+        The ExecuTorch schema is:
+            aten.avg_pool2d(
+                Tensor self,
+                int[2] kernel_size,
+                int[2] stride=[],
+                int[2] padding=0,
+                bool ceil_mode=False
+                bool count_include_pad=True,
+                int? divisor_override=None
+            )
+        """
         self.assert_convertible(node)
 
         kernel_size = node.args[1]

backends/nxp/nxp_backend.py

@@ -15,6 +15,9 @@
 import numpy as np
 import torch
 
+from executorch.backends.nxp.backend.custom_delegation_options import (
+    CustomDelegationOptions,
+)
 from executorch.backends.nxp.backend.data_format import DataFormat
 from executorch.backends.nxp.backend.edge_program_converter import (
     EdgeProgramToIRConverter,
@@ -229,6 +232,9 @@ def preprocess(  # noqa C901
                 edge_program,
                 neutron_target_spec=NeutronTargetSpec(target),
                 conversion_config=conversion_config,
+                custom_delegation_options=CustomDelegationOptions(
+                    use_new_flow_neutron_c=use_new_flow_neutron_c
+                ),
             )
 
             neutron_model = NeutronConverterManager(dump_kernel_selection_code).convert(

backends/nxp/quantizer/utils.py

@@ -1,5 +1,5 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
-# Copyright 2024-2025 NXP
+# Copyright 2024-2026 NXP
 # All rights reserved.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -10,7 +10,7 @@
 import itertools
 from collections import OrderedDict
 from collections.abc import Iterable
-from typing import Any, Dict, List, Tuple, Type
+from typing import Any, Callable, Dict, List, Tuple, Type
 
 import torch
 from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_linear_pass import (
@@ -30,8 +30,10 @@
     check_subgraphs_connected,
     SourcePartition,
 )
+
 from torchao.quantization.pt2e import (
     move_exported_model_to_eval,
+    move_exported_model_to_train,
     ObserverOrFakeQuantize,
 )
 from torchao.quantization.pt2e.quantize_pt2e import (
@@ -176,16 +178,17 @@ def calibrate_and_quantize(
     calibration_inputs: Iterable[tuple[torch.Tensor, ...]],
     quantizer: Quantizer,
     is_qat: bool = False,
+    train_fn: Callable[[torch.fx.GraphModule], None] | None = None,
 ) -> fx.GraphModule:
     """Quantize the provided model.
 
     :param model: Aten model (or it's GraphModule representation) to quantize.
-    :param calibration_inputs: Either a tuple of calibration input tensors where each element corresponds to a model
-                                input. Or an iterator over such tuples.
+    :param calibration_inputs: An iterator over tuples of calibration input tensors where each tensor corresponds to a
+                                model input.
     :param quantizer: Quantizer to use.
     :param is_qat: Whether quantization is done using Quantization Aware Training (QAT) or not.
                     Note: In QAT mode, training is not performed. Only calibration (in eval mode) is done.
-
+    :param train_fn: Optional training function to be called during QAT.
     :return: Quantized GraphModule.
     """
 
@@ -195,12 +198,20 @@ def calibrate_and_quantize(
     if is_qat:
         m = prepare_qat_pt2e(model, quantizer)
         m = AddSimulatedLinearBatchNormFusionQATPass()(m).graph_module
+
+        if train_fn:
+            m = move_exported_model_to_train(m)
+            train_fn(m)
+
         m = move_exported_model_to_eval(m)
+        m = RemoveSimulatedLinearBatchNormFusionQATPass()(m).graph_module
+        m = FuseBatchNormWithLinearPass()(m).graph_module
     else:
         m = prepare_pt2e(model, quantizer)
 
-    for data in calibration_inputs:
-        m(*data)
+    if not is_qat or (is_qat and not train_fn):
+        for data in calibration_inputs:
+            m(*data)
 
     if is_qat:
         m = RemoveSimulatedLinearBatchNormFusionQATPass()(m).graph_module

backends/nxp/tests_models/config.py → backends/nxp/tests/config.py

@@ -21,7 +21,7 @@
 
 NSYS_PATH = pathlib.Path(shutil.which("nsys"))
 NSYS_CONFIG_PATH = os.path.join(
-    PROJECT_DIR, "backends", "nxp", "tests_models", "neutron-imxrt700.ini"
+    PROJECT_DIR, "backends", "nxp", "tests", "neutron-imxrt700.ini"
 )
 NSYS_FIRMWARE_PATH = os.path.join(
     os.path.dirname(eiq_neutron_sdk.__file__),

backends/nxp/tests_models/config_importer.py → backends/nxp/tests/config_importer.py

@@ -12,6 +12,6 @@
 
     logger.debug("Importing from executorch-integration")
 except ImportError:
-    import executorch.backends.nxp.tests_models.config as test_config  # noqa F401
+    import executorch.backends.nxp.tests.config as test_config  # noqa F401
 
     logger.debug("Importing from executorch")

backends/nxp/tests_models/conftest.py → backends/nxp/tests/conftest.py

@@ -8,7 +8,7 @@
 import pathlib
 import shutil
 
-from executorch.backends.nxp.tests_models.outputs_dir_importer import outputs_dir
+from executorch.backends.nxp.tests.outputs_dir_importer import outputs_dir
 
 
 def pytest_addoption(parser):

backends/nxp/tests_models/dataset_creator.py → backends/nxp/tests/dataset_creator.py

@@ -14,8 +14,8 @@
 import numpy as np
 import torch
 from executorch.backends.nxp.backend.ir.converter.conversion import translator
-from executorch.backends.nxp.tests_models.calibration_dataset import CalibrationDataset
-from executorch.backends.nxp.tests_models.model_input_spec import ModelInputSpec
+from executorch.backends.nxp.tests.calibration_dataset import CalibrationDataset
+from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec
 from torch import Tensor