NXP backend: handle tests of mlperf tiny keyword spotting

backends/nxp/tests/generic_tests/test_aot_example.py

@@ -249,3 +249,69 @@ def test_aot_example__mlperf_tiny_ic__profiling():
     with _cleanup_generated_files(pte_file, etrecord_file):
         result = _run_compile(cmd)
         _assert_profiling(result, pte_file, etrecord_file)
+
+
+def test_aot_example__mlperf_tiny_kws():
+    """Test that the MLPerf Tiny keyword spotting model (DS-CNN) can be lowered to Neutron backend via
+    `aot_neutron_compile.py` and all ops are delegated."""
+
+    # Run the compilation script as a module (like run_aot_example.sh does).
+    # The calibration data of this model is generated randomly, so no dataset download is needed.
+    cmd = [
+        sys.executable,
+        "-m",
+        "examples.nxp.aot_neutron_compile",
+        "--model_name",
+        "mlperf_tiny_keyword_spotting",
+        "--delegate",
+        "--quantize",
+        "--target",
+        "imxrt700",
+        "--use_random_dataset",
+    ]
+
+    # Output file will be created in executorch_root
+    pte_file = Path(
+        os.path.join(EXECUTORCH_ROOT, "mlperf_tiny_keyword_spotting_nxp_delegate.pte")
+    )
+
+    with _cleanup_generated_files(pte_file):
+        result = _run_compile(cmd)
+        _assert_delegation(result, pte_file)
+
+
+def test_aot_example__mlperf_tiny_kws__profiling():
+    """Test that the MLPerf Tiny keyword spotting model (DS-CNN) can be lowered to Neutron backend via
+    `aot_neutron_compile.py` and all ops are delegated."""
+
+    # Run the compilation script as a module (like run_aot_example.sh does)
+    cmd = [
+        sys.executable,
+        "-m",
+        "examples.nxp.aot_neutron_compile",
+        "--model_name",
+        "mlperf_tiny_keyword_spotting",
+        "--delegate",
+        "--quantize",
+        "--target",
+        "imxrt700",
+        "--remove-quant-io-ops",
+        "--use_profiling",  # Generate profilable model and create ETRecord
+        "--use_random_dataset",  # Avoid downloading the dataset.
+    ]
+
+    # Output files will be created in executorch_root.
+    pte_file = Path(
+        os.path.join(
+            EXECUTORCH_ROOT, "mlperf_tiny_keyword_spotting_nxp_delegate_profile.pte"
+        )
+    )
+    etrecord_file = Path(
+        os.path.join(
+            EXECUTORCH_ROOT, "etrecord", "mlperf_tiny_keyword_spotting_etrecord.bin"
+        )
+    )
+
+    with _cleanup_generated_files(pte_file, etrecord_file):
+        result = _run_compile(cmd)
+        _assert_profiling(result, pte_file, etrecord_file)

backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py

@@ -0,0 +1,128 @@
+# 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 functools import partial
+
+import numpy as np
+import torch
+from executorch.backends.nxp.tests.nsys_testing import ReferenceModel
+from executorch.backends.nxp.tests.dataset_creator import (
+    FromCalibrationDataDatasetCreator,
+)
+from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec
+from executorch.backends.nxp.tests.graph_verifier import BaseGraphVerifier
+from executorch.backends.nxp.tests.model_output_comparator import (
+    NumericalStatsOutputComparator,
+)
+
+from executorch.backends.nxp.tests.nsys_testing import (
+    lower_run_compare,
+)
+from executorch.backends.nxp.tests.use_qat import *  # noqa F403
+import pytest
+from executorch.examples.nxp.models.mlperf_tiny.keyword_spotting.mlperf_tiny_keyword_spotting import (
+    MLPerfTinyKeywordSpotting,
+)
+
+BOUNDS_MSE = {
+    "PTQ": {
+        "channels-last": 9.5e-3,
+        "channels-first": 9.5e-3,
+    },
+    "QAT": {
+        "channels-last": 9.5e-3,
+        "channels-first": 9.5e-3,
+    },
+}
+
+
+@pytest.fixture(autouse=True)
+def reseed_model_per_test_run():
+    torch.manual_seed(23)
+    np.random.seed(23)
+
+
+@pytest.mark.parametrize("channels_last", [True])
+def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last):
+    channels_last = True
+    use_qat = False
+    # 5 samples per class
+    num_samples = 60
+
+    kws = MLPerfTinyKeywordSpotting(num_samples=num_samples, use_random_dataset=True)
+    model = kws.get_eager_model()
+    dataset = kws.dataset
+    labels = kws.labels
+
+    dataset_creator = FromCalibrationDataDatasetCreator(
+        dataset, num_examples=num_samples, idx_to_label=labels
+    )
+
+    input_spec = ModelInputSpec(kws.input_shape)
+    if channels_last:
+        model.to(memory_format=torch.channels_last)
+        input_spec.dim_order = torch.channels_last
+
+    bounds_key_1 = "QAT" if use_qat else "PTQ"
+    bounds_key_2 = "channels-last" if channels_last else "channels-first"
+    mse = BOUNDS_MSE[bounds_key_1][bounds_key_2]
+    comparator = NumericalStatsOutputComparator(
+        max_mse_error=mse, is_classification_task=True
+    )
+    model_verifier = BaseGraphVerifier(1, [])
+    train_fn = (
+        partial(kws.train_model_fn, channels_last=channels_last)
+        if use_qat
+        else None
+    )
+
+    ref_model = (
+        ReferenceModel.QUANTIZED_EDGE_PYTHON
+        if channels_last
+        else ReferenceModel.QUANTIZED_EXECUTORCH_CPP
+    )
+
+    lower_run_compare(
+        model,
+        [input_spec],
+        model_verifier,
+        request,
+        dataset_creator=dataset_creator,
+        output_comparator=comparator,
+        mocker=mocker,
+        reference_model=ref_model,
+        use_qat=use_qat,
+        train_fn=train_fn,
+    )
+
+
+# @pytest.mark.xfail(reason="EIEX-512", strict=True)
+# def test_mlperf_tiny_kws_ptq_qat_equivalence(request):
+#     # 5 samples per class
+#     num_samples = 60
+
+#     kws = MLPerfTinyKeywordSpotting(num_samples=num_samples, use_random_dataset=True)
+
+#     model = kws.get_eager_model()
+#     dataset = kws.dataset
+#     labels = kws.labels
+
+#     dataset_creator = FromCalibrationDataDatasetCreator(
+#         dataset, num_examples=num_samples, idx_to_label=labels
+#     )
+#     comparator = ClassificationAccuracyOutputComparator(class_dict=labels)
+
+#     input_spec = ModelInputSpec(kws.input_shape)
+#     model_verifier = BaseGraphVerifier(1, [])
+
+#     lower_run_compare_ptq_qat(
+#         model,
+#         [input_spec],
+#         model_verifier,
+#         request,
+#         train_fn=kws.train_model_fn,
+#         dataset_creator=dataset_creator,
+#         output_comparator=comparator,
+#     )

examples/nxp/aot_neutron_compile.py

@@ -45,6 +45,9 @@
 from executorch.examples.nxp.models.mlperf_tiny.image_classification.mlperf_tiny_image_classification import (
     MLPerfTinyImageClassification,
 )
+from executorch.examples.nxp.models.mlperf_tiny.keyword_spotting.mlperf_tiny_keyword_spotting import (
+    MLPerfTinyKeywordSpotting,
+)
 from executorch.examples.nxp.models.mobilenet_v2 import MobilenetV2
 from executorch.exir import (
     EdgeCompileConfig,
@@ -64,6 +67,7 @@
     "cifar10": CifarNet,
     "mobilenetv2": MobilenetV2,
     "mlperf_tiny_image_classification": MLPerfTinyImageClassification,
+    "mlperf_tiny_keyword_spotting": MLPerfTinyKeywordSpotting,
 }
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
@@ -118,7 +122,7 @@ def _get_model_info_from_name(
                 )
             model_cls_inst = model_cls()
 
-        elif model_cls is MLPerfTinyImageClassification:
+        elif model_cls in (MLPerfTinyImageClassification, MLPerfTinyKeywordSpotting):
             model_cls_inst = model_cls(
                 dataset_path=dataset_path, use_random_dataset=use_random_dataset
             )
@@ -316,7 +320,8 @@ def _get_arg_parser():
         quantizer = NeutronQuantizer(neutron_target_spec, is_qat=args.use_qat)
         if args.use_qat:
             if not isinstance(
-                model_cls_inst, (CifarNet, MLPerfTinyImageClassification)
+                model_cls_inst,
+                (CifarNet, MLPerfTinyImageClassification, MLPerfTinyKeywordSpotting),
             ):
                 raise ValueError(
                     f"QAT training is not supported for model '{args.model_name}'"

examples/nxp/models/mlperf_tiny/keyword_spotting/__init__.py

@@ -0,0 +1,4 @@
+# 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.

examples/nxp/models/mlperf_tiny/keyword_spotting/mlperf_tiny_keyword_spotting.py

@@ -0,0 +1,79 @@
+# 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 pathlib import Path
+
+import torch
+
+from executorch.backends.nxp.tests.calibration_dataset import (
+    CalibrationDataset,
+    RandomCalibrationDataset,
+)
+
+from executorch.examples.models.mlperf_tiny import DSCNNKWS
+from executorch.examples.nxp.models.mlperf_tiny.mlperf_tiny_model import MLPerfTinyModel
+from torch.utils.data import Dataset
+
+
+log = logging.getLogger(__name__)
+
+INPUT_SHAPE = (1, 1, 49, 10)
+IDX_TO_LABEL = {
+    0: "Down",
+    1: "Go",
+    2: "Left",
+    3: "No",
+    4: "Off",
+    5: "On",
+    6: "Right",
+    7: "Stop",
+    8: "Up",
+    9: "Yes",
+    10: "Silence",
+    11: "Unknown",
+}
+
+
+class MLPerfTinyKeywordSpotting(MLPerfTinyModel):
+    """MLPerf Tiny keyword spotting model (DS-CNN)."""
+
+    def __init__(
+        self,
+        num_samples: int = 200,
+        dataset_path: Path | str | None = None,
+        use_random_dataset: bool = False,
+    ):
+        self._num_samples = num_samples
+        self._use_random_dataset = use_random_dataset
+        self._dataset_path = dataset_path
+
+        super().__init__()
+
+    @property
+    def input_shape(self):
+        return INPUT_SHAPE
+
+    @property
+    def labels(self):
+        return IDX_TO_LABEL
+
+    def _init_dataset(self) -> Dataset:
+        if self._use_random_dataset:
+            num_classes = len(self.labels)
+            sample_shape = tuple(self.input_shape)[1:]
+            return RandomCalibrationDataset(
+                self._num_samples, sample_shape, num_classes
+            )
+        else:
+            if self._dataset_path is None:
+                raise ValueError(
+                    "Path to dataset data cannot be empty. If you want to use random data, set `use_random_dataset = True`"
+                )
+            return CalibrationDataset(self._dataset_path)
+
+    def _init_eager_model(self) -> torch.nn.Module:
+        num_classes = len(self.labels)
+        return DSCNNKWS(num_classes)

examples/nxp/models/mlperf_tiny/mlperf_tiny_model.py

@@ -4,12 +4,17 @@
 # LICENSE file in the root directory of this source tree.
 
 import itertools
+import logging
 from abc import abstractmethod
 from typing import Iterator
 
 import torch
 from executorch.examples.models import model_base
 from torch.utils.data import DataLoader, Dataset
+from torchao.quantization.pt2e import disable_observer
+from tqdm import tqdm
+
+log = logging.getLogger(__name__)
 
 
 class MLPerfTinyModel(model_base.EagerModelBase):
@@ -79,3 +84,34 @@ def get_eager_model(self):
 
     def get_example_inputs(self) -> tuple[torch.Tensor]:
         return (torch.randn(self.input_shape, dtype=torch.float32),)
+
+    def train_model_fn(self, model, num_epochs=15, batch_size=64, channels_last=False):
+        torch.manual_seed(42)
+        torch.use_deterministic_algorithms(True)
+
+        optimizer = torch.optim.Adam(
+            params=model.parameters(),
+            lr=5e-6,
+            eps=1e-7,
+            weight_decay=1e-4,
+        )
+        loss_fn = torch.nn.CrossEntropyLoss()
+
+        log.warning("Starting training...")
+
+        data = self.get_qat_train_inputs(batch_size=batch_size)
+        for nepoch in range(num_epochs):
+            for samples, labels in tqdm(data):
+                if channels_last:
+                    samples = samples.to(memory_format=torch.channels_last)
+
+                optimizer.zero_grad()
+                outputs = model(samples)
+                loss = loss_fn(outputs, labels)
+                loss.backward()
+                optimizer.step()
+
+            if nepoch >= num_epochs / 3:
+                model.apply(disable_observer)
+
+        return model