NXP backend: Unify testing infrastructure. (pytorch#19112)

### Summary
Unify infrastructure for tests comparing edge dialect <-> Neutron IR
with tests comparing lowered edge programs <-> lowered Neutron delegated
programs.

### Test plan
Tested by all existing tests.


cc @robert-kalmar @JakeStevens @digantdesai

Author: MartinPavella

.github/workflows/pull.yml

@@ -1235,26 +1235,26 @@ jobs:
 
         # Build and install Executorch
         PYTHON_EXECUTABLE=python \
-        CMAKE_ARGS="-DEXECUTORCH_BUILD_NXP_NEUTRON=ON -DEXECUTORCH_BUILD_NXP_NEUTRON_RUNNER=ON " \
-        .ci/scripts/setup-linux.sh --build-tool "cmake"
+        CMAKE_ARGS="-DEXECUTORCH_BUILD_NXP_NEUTRON=ON -DEXECUTORCH_BUILD_NXP_NEUTRON_RUNNER=ON \
+        -DEXECUTORCH_BUILD_KERNELS_PORTABLE=ON -DEXECUTORCH_BUILD_PYBIND=ON -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
+        -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \
+        -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON -DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=ON" \
+        .ci/scripts/setup-linux.sh --build-tool "cmake"  --editable true
 
         # Install test requirements
         pip install -r backends/nxp/requirements-tests-pypi.txt
         PYTHON_EXECUTABLE=python bash examples/nxp/setup.sh
 
-        # Run pytest
-        PYTHON_EXECUTABLE=python bash backends/nxp/run_unittests.sh
-
         # Run aot examples:
         PYTHON_EXECUTABLE=python bash examples/nxp/run_aot_example.sh cifar10
         PYTHON_EXECUTABLE=python bash examples/nxp/run_aot_example.sh mobilenetv2
 
         # Run e2e example with Simulator:
         PYTHON_EXECUTABLE=python bash examples/nxp/run.sh cifar10
 
-        # Run lightweight model tests:
-        PYTHON_EXECUTABLE=python pytest -c /dev/null backends/nxp/tests_models/ \
-        --nxp_runner_path "./examples/nxp/executor_runner/build/nxp_executor_runner"
+        # Run unit tests:
+        PYTHON_EXECUTABLE=python NXP_RUNNER_PATH="./examples/nxp/executor_runner/build/nxp_executor_runner" \
+        bash backends/nxp/run_unittests.sh
 
   test-samsung-quantmodels-linux:
     name: test-samsung-quantmodels-linux

backends/nxp/__init__.py

@@ -1,5 +1,5 @@
 # Copyright 2023 Martin Pavella
-# Copyright 2023-2025 NXP
+# Copyright 2023-2026 NXP
 #
 # License: MIT
 # See the LICENSE_MIT for more details.
@@ -547,8 +547,8 @@ def convert_data_type(torch_type: torch.TensorType) -> TensorType:
         )
 
 
-def torch_type_to_numpy_type(torch_type: torch.TensorType) -> np.ScalarType:
-    """Convert Torch DataType to NeutronIR TensorType"""
+def torch_type_to_numpy_type(torch_type: torch.TensorType) -> np.dtype:
+    """Convert Torch data type to Numpy data type."""
 
     if torch_type == torch.float32:
         return np.dtype(np.float32)

backends/nxp/backend/ir/converter/conversion/translator.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/quantizer/utils.py

@@ -15,13 +15,19 @@
 # TODO(Robert Kalmar) In accordance with the "TODO(dbort): Prune /test[s]/ dirs, /third-party/ dirs" in pyproject.toml,
 #  once the test folders are not installed we can derive the path from current file location: `pathlib.Path(__file__)`
 PROJECT_DIR = os.environ.get("PROJECT_DIR")
-assert PROJECT_DIR and os.path.exists(PROJECT_DIR)
+if not PROJECT_DIR:
+    # Auto-detect: The PROJECT_DIR env variable is set by the conftest.py in backends.nxp.tests.conftest. But unittests
+    #               don't use the conftest, so this variable is not set -> set it manually here in that case.
+    PROJECT_DIR = str(pathlib.Path(__file__).parent.parent.parent.parent)
+assert os.path.exists(
+    PROJECT_DIR
+), f"Invalid PROJECT_DIR env variable: `{PROJECT_DIR}`."
 
 OUTPUTS_DIR = pathlib.Path(os.getcwd()) / ".outputs"
 
 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__),
@@ -34,4 +40,17 @@
 # The NXP_RUNNER_PATH env variable is either defined by pytest when using the CLI argument --nxp_executor_path or
 # a standard environment variable.
 NEUTRON_TEST_PATH = os.environ.get("NXP_RUNNER_PATH")
-assert NEUTRON_TEST_PATH and os.path.exists(NEUTRON_TEST_PATH)
+if not NEUTRON_TEST_PATH:
+    # Auto-detect: The NXP_RUNNER_PATH env variable is set by the conftest.py in backends.nxp.tests.conftest. But
+    #               unittests don't use the conftest, so this variable is not set -> set it manually here in that case.
+    NEUTRON_TEST_PATH = (
+        pathlib.Path(PROJECT_DIR)
+        / "examples"
+        / "nxp"
+        / "executor_runner"
+        / "build"
+        / "nxp_executor_runner"
+    )
+assert os.path.exists(
+    NEUTRON_TEST_PATH
+), f"Invalid NXP_RUNNER_PATH env variable: `{NEUTRON_TEST_PATH}`."

backends/nxp/tests/__init__.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")

…/nxp/tests_models/calibration_dataset.py backends/nxp/tests/calibration_dataset.pybackends/nxp/tests_models/calibration_dataset.py renamed to backends/nxp/tests/calibration_dataset.py backends/nxp/tests_models/calibration_dataset.py renamed to backends/nxp/tests/calibration_dataset.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/config.py backends/nxp/tests/config.pybackends/nxp/tests_models/config.py renamed to backends/nxp/tests/config.py

@@ -14,8 +14,11 @@
 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.backend.ir.converter.conversion.translator import (
+    torch_type_to_numpy_type,
+)
+from executorch.backends.nxp.tests.calibration_dataset import CalibrationDataset
+from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec
 from torch import Tensor
 
 
@@ -100,10 +103,15 @@ def _gen_samples(
                     case _:
                         raise ValueError(f"Unsupported dim_order: {spec.dim_order}")
 
-                sample_vector = rng.random(np.prod(shape), spec.type).reshape(shape)
-                sample_vector.tofile(
-                    os.path.join(sample_dir, f"{str(spec_idx).zfill(2)}.bin")
+                sample_vector = rng.random(
+                    np.prod(shape), torch_type_to_numpy_type(spec.dtype)
+                ).reshape(shape)
+                file_name = (
+                    f"{str(spec_idx).zfill(2)}.bin"
+                    if len(input_spec) > 1
+                    else f"{str(idx).zfill(4)}.bin"
                 )
+                sample_vector.tofile(os.path.join(sample_dir, file_name))
 
 
 class CopyDatasetCreator(DatasetCreator):
@@ -132,9 +140,9 @@ def generate_samples(self, dataset_dir, input_spec) -> tuple[str, str]:
 
             if input_spec[0].dim_order == torch.channels_last:
                 # Permute the sample to channels last and store it in the testing dataset.
-                tensor = np.fromfile(sample_path, dtype=input_spec[0].type).reshape(
-                    input_spec[0].shape
-                )
+                tensor = np.fromfile(
+                    sample_path, dtype=torch_type_to_numpy_type(input_spec[0].dtype)
+                ).reshape(input_spec[0].shape)
 
                 if (
                     list(tensor.shape) == list(input_spec[0].shape)