NXP backend: Fix adding observers on non-float nodes and Fix incorrect turn of an output buffer into model parameter (#23890)

### Summary
Fixing errors observed by executorch-fuzzer. 

During calibration, the quantizer may inject operations that are
incompatible with the graph's data types; for example, it may invoke
torch.histc on int64 tensors, which the operation does not support. In
other cases, the conversion pass may incorrectly turn an output buffer
into amodel parameter while the graph continues to mutate that
buffer,causing verification to fail.

### Test plan
Unit test introduced

Signed-off-by: Robert Kalmar <robert.kalmar@nxp.com>

Author: robert-kalmar

backends/nxp/aten_passes/remove_nodes_with_known_outputs.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.
@@ -70,6 +70,25 @@ def _is_getitem(node_: Node) -> bool:
         users = list(node.users.keys())
         return all(_is_getitem(user) for user in users)
 
+    @staticmethod
+    def node_is_used_as_out_kwarg(node: Node) -> bool:
+        """Return `True` if `node` is passed as the `out=` kwarg of any of its users.
+
+        Such a node (typically a `torch.empty(...)` placeholder buffer feeding a `*.out`
+        overload) is about to be mutated in place by its consumer, so its "known" data
+        (e.g. uninitialized garbage from `empty.memory_format`) is not actually the value
+        observed at runtime, and it must not be folded into a static `Parameter`: `torch.export`
+        would then treat the mutated buffer as a learned weight, which breaks the exported
+        program's parameter/signature verification later in the pipeline.
+        """
+        for user in node.users:
+            if user.kwargs.get("out") is node:
+                return True
+            out_args = user.kwargs.get("out")
+            if isinstance(out_args, (list, tuple)) and node in out_args:
+                return True
+        return False
+
     def replace_node_with_static_data(self, node: Node, static_data: Parameter):
         """Remove the given `node` from the graph and replace it with a parameter node containing the `static_data`."""
         # Generate a unique name for the new static parameter.
@@ -104,7 +123,6 @@ def replace_following_getitem_nodes_with_static_data(
                              after the `root_node`.
         :return: `True` if the replacement was successfully executed. `False` otherwise.
         """
-
         if not self.node_is_followed_only_by_getitem_nodes(root_node):
             return False  # Unexpected case.
 
@@ -158,6 +176,12 @@ def call(self, module: GraphModule) -> PassResult:
             if node.op != "call_function":
                 continue  # Not a compute operator.
 
+            if self.node_is_used_as_out_kwarg(node):
+                # `node` is the `out=` buffer of a mutating `*.out` overload. Its "known"
+                #  output data is not what gets observed at runtime (it is about to be
+                #  overwritten), so it must not be folded into a static `Parameter`.
+                continue
+
             # Try to access the static data for the inputs of the node.
             args_with_data = self.replace_nodes_in_list_with_their_data(node.args)
 

backends/nxp/quantizer/neutron_quantizer.py

@@ -100,6 +100,32 @@
 from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY
 
 
+#: Dtypes accepted by `quantized_decomposed.quantize_per_tensor` (inserted by `convert_pt2e`)
+_QUANTIZABLE_FLOAT_DTYPES = (torch.float32, torch.float16, torch.bfloat16)
+
+
+def _is_quantizable_node(node: fx.Node) -> bool:
+    """Returns whether node is quantizable float type, per `quantized_decomposed.quantize_per_tensor` torch
+     operator.
+
+    Helper function to avoid placing Observers or FakeQuant on non-floating nodes.
+    Nodes without a `"val"` meta entry (e.g. linear_weight_bias) are treated as quantizable.
+    """
+    val = node.meta.get("val", None)
+    if val is None or not hasattr(val, "dtype"):
+        return True
+    return val.dtype in _QUANTIZABLE_FLOAT_DTYPES
+
+
+def _dtype_safe_qspec(
+    spec: QuantizationSpec | None, node: fx.Node
+) -> QuantizationSpec | None:
+    """Suppress quantization annotation for a non-floating-point node inputs."""
+    if _is_quantizable_node(node):
+        return spec
+    return None
+
+
 class NeutronAtenQuantizer(Quantizer):
     def __init__(
         self, pattern: QuantizationPattern, quantization_config: QuantizationConfig
@@ -109,6 +135,7 @@ def __init__(
         self.quantization_config = quantization_config
 
     def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
+
         fused_partitions = find_sequential_partitions_aten(
             model,
             self.pattern.partition_types(),
@@ -138,12 +165,15 @@ def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
                 continue
 
             for output, *custom_spec in anchors.output:
+                spec = custom_spec[0] if custom_spec else output_act_qspec
+                spec = _dtype_safe_qspec(spec, output)
                 # pyre-ignore[16]: no attribute
-                output.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(
-                    # pyre-ignore[6]: incompatible parameter type
-                    output_qspec=(custom_spec[0] if custom_spec else output_act_qspec),
-                    _annotated=True,
-                )
+                if spec is not None:
+                    output.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(
+                        # pyre-ignore[6]: incompatible parameter type
+                        output_qspec=spec,
+                        _annotated=True,
+                    )
 
             def annotate_inputs(
                 inputs: (
@@ -165,9 +195,9 @@ def annotate_inputs(
                         # pyre-ignore[16]: no attribute
                         else node.args[args_idx.idx][args_idx.inner_idx]
                     )
-                    annotation.input_qspec_map[arg] = (
-                        custom_spec[0] if custom_spec else spec
-                    )
+                    arg_spec = custom_spec[0] if custom_spec else spec
+                    arg_spec = _dtype_safe_qspec(arg_spec, arg)
+                    annotation.input_qspec_map[arg] = arg_spec
                     # pyre-ignore[16]: no attribute
                     node.meta[Q_ANNOTATION_KEY] = annotation
 
@@ -399,8 +429,10 @@ def _annotate_inputs(self, model: fx.GraphModule):
                 continue
 
             if node.op == "placeholder" and len(node.users) > 0:
-                annotate_output_qspec(node, act_qspec(self.is_qat))
-                self._mark_input_node_as_annotated(node)
+                spec = _dtype_safe_qspec(act_qspec(self.is_qat), node)
+                if spec is not None:
+                    annotate_output_qspec(node, spec)
+                    self._mark_input_node_as_annotated(node)
 
     def validate(self, model: torch.fx.GraphModule) -> None:
         return super().validate(model)

backends/nxp/quantizer/utils.py

@@ -32,8 +32,6 @@
 )
 
 from torchao.quantization.pt2e import (
-    HistogramObserver,
-    MinMaxObserver,
     move_exported_model_to_eval,
     move_exported_model_to_train,
     ObserverOrFakeQuantize,
@@ -202,42 +200,6 @@ def find_sequential_partitions_aten(
     return fused_partitions
 
 
-def _replace_histogram_observers_for_integer_inputs(
-    m: torch.fx.GraphModule,
-) -> None:
-    """Replace HistogramObserver with MinMaxObserver for observer nodes whose
-    input tensor has a non-floating-point dtype.
-
-    HistogramObserver calls torch.histc internally, which raises
-    NotImplementedError for integer and bool dtypes.  MinMaxObserver only
-    tracks min/max and handles all dtypes, producing equally valid
-    quantization ranges for non-float inputs (e.g. embedding indices, bool
-    masks).
-    """
-    for node in m.graph.nodes:
-        if node.op != "call_module":
-            continue
-        obs = getattr(m, node.target, None)
-        if not isinstance(obs, HistogramObserver):
-            continue
-        input_node = node.args[0] if node.args else None
-        if input_node is None:
-            continue
-        val = input_node.meta.get("val", None)
-        if val is not None and not val.is_floating_point():
-            setattr(
-                m,
-                node.target,
-                MinMaxObserver(
-                    dtype=obs.dtype,
-                    qscheme=obs.qscheme,
-                    quant_min=obs.quant_min,
-                    quant_max=obs.quant_max,
-                    eps=obs.eps,
-                ),
-            )
-
-
 def calibrate_and_quantize(
     model: ExportedProgram | fx.GraphModule,
     calibration_inputs: Iterable[tuple[torch.Tensor, ...]],
@@ -266,13 +228,6 @@ def calibrate_and_quantize(
     else:
         m = prepare_pt2e(model, quantizer)
 
-    # Swap HistogramObserver -> MinMaxObserver for any observer whose input has
-    # a non-floating-point dtype.  Safe for both PTQ and QAT: in QAT the graph
-    # contains FakeQuantize nodes which are not HistogramObserver instances, so
-    # the helper is effectively a no-op there, but it protects against the edge
-    # case where a QAT graph does contain a HistogramObserver.
-    _replace_histogram_observers_for_integer_inputs(m)
-
     if is_qat:
         if train_fn:
             m = move_exported_model_to_train(m)

backends/nxp/recipes/nxp_recipe_provider.py

@@ -38,9 +38,6 @@
     default_preserve_ops,
     generate_neutron_compile_spec,
 )
-from executorch.backends.nxp.quantizer.utils import (
-    _replace_histogram_observers_for_integer_inputs,
-)
 from executorch.backends.nxp.recipes.nxp_recipe_types import NXP_BACKEND, NXPRecipeType
 from executorch.backends.nxp.tests.executorch_pipeline import (
     get_default_quantizer,
@@ -132,7 +129,6 @@ class NeutronRecipeConfig:
 
 
 class NXPRecipeProvider(BackendRecipeProvider):
-
     @property
     def backend_name(self) -> str:
         return NXP_BACKEND
@@ -245,14 +241,9 @@ def _wrapped(m):
     return _wrapped
 
 
-def _histogram_observer_fix_pass(m):
-    """Callable(GraphModule) -> GraphModule that replaces HistogramObserver for integer inputs."""
-    _replace_histogram_observers_for_integer_inputs(m)
-    return m
-
-
 # ---------------------------------------------------------------------------
 # Module-level builder helpers
+
 # ---------------------------------------------------------------------------
 
 
@@ -269,13 +260,11 @@ def _build_quantization_recipe(
     """
     _quantizer = rc.get_quantizer_fn()
 
-    # post_prepare_passes: always fix HistogramObserver for non-float inputs.
-    # For QAT, also insert the simulated linear-BN fusion before training so
-    # fake-quantize nodes see fused weights during the training loop.
+    # post_prepare_passes: for QAT, insert the simulated linear-BN fusion before
+    # training so fake-quantize nodes see fused weights during the training loop.
     post_prepare: list[Callable] = []
     if is_qat:
         post_prepare.append(_wrap_exir_pass(AddSimulatedLinearBatchNormFusionQATPass))
-    post_prepare.append(_histogram_observer_fix_pass)
 
     if is_qat:
         # pre_convert_passes: tear down the simulated fusion and fold BN into

backends/nxp/tests/generic_tests/test_quantizer.py

@@ -2,6 +2,7 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+import functools
 
 # Tests for NeutronQuantizer.
 
@@ -39,6 +40,7 @@
 )
 from executorch.backends.nxp.tests.executors import (
     convert_run_compare,
+    graph_contains_any,
     graph_contains_any_of_ops,
     tflite,
     ToChannelFirstPreprocess,
@@ -760,3 +762,44 @@ def is_conv(node):
         assert len(graph_nodes) == 21
     else:
         assert len(graph_nodes) == 15
+
+
+def test_quantizer__non_float_input_does_not_crash_prepare_or_convert():
+    """A non-floating-point input (e.g. int64) must not crash the
+    prepare_pt2e/convert_pt2e.
+
+    Uses prepare_pt2e/convert_pt2e directly (not calibrate_and_quantize test
+    wrapper), matching how any external PT2E caller might use NeutronQuantizer.
+    """
+
+    class PowModule(torch.nn.Module):
+        def forward(self, x, y):
+            return torch.ops.aten.pow.Tensor_Tensor(x, y)
+
+    x = torch.randint(-4, 5, (2,), dtype=torch.int64)
+    y = torch.randn((2,), dtype=torch.float32)
+
+    exported_model = export(PowModule(), (x, y), strict=True)
+    quantizer = NeutronQuantizer(neutron_target_spec)
+
+    prepared_model = prepare_pt2e(exported_model.module(), quantizer)
+    prepared_model(x, y)  # calibration must not raise (torch.histc on int64)
+
+    quantized_model = convert_pt2e(
+        prepared_model
+    )  # must not raise (quantize_per_tensor dtype check)
+    quantized_model(x, y)  # sanity: the converted graph still runs
+
+    def quant_node_on(n, name):
+        return (
+            hasattr(n, "target")
+            and n.target is torch.ops.quantized_decomposed.quantize_per_tensor.default
+            and n.all_input_nodes[0].name == name
+        )
+
+    assert graph_contains_any(
+        quantized_model.graph, condition=functools.partial(quant_node_on, name="y")
+    )
+    assert not graph_contains_any(
+        quantized_model.graph, condition=functools.partial(quant_node_on, name="x")
+    )

backends/nxp/tests/generic_tests/test_recipe_export.py

@@ -24,7 +24,6 @@
 )
 from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
 from executorch.backends.nxp.recipes.nxp_recipe_provider import (
-    _histogram_observer_fix_pass,
     NEUTRON_RECIPE_CONFIG_KEY,
     NeutronRecipeConfig,
     NXPRecipeProvider,
@@ -609,14 +608,6 @@ def test__post_prepare_passes_start_with_add_bn_fusion(self, recipe_type):
             AddSimulatedLinearBatchNormFusionQATPass.__name__ in first_pass.__qualname__
         )
 
-    def test__post_prepare_passes_end_with_histogram_observer_fix(self, recipe_type):
-        """QAT post_prepare_passes: last pass is _histogram_observer_fix_pass."""
-        recipe = self._get_qat_recipe(recipe_type)
-        qr = recipe.quantization_recipe
-        assert qr.post_prepare_passes is not None
-        last_pass = qr.post_prepare_passes[-1]
-        assert last_pass is _histogram_observer_fix_pass
-
     def test__pre_convert_passes_include_remove_bn_fusion_and_fold(self, recipe_type):
         """QAT pre_convert_passes: contains RemoveSimulatedLinearBatchNormFusionQATPass
         followed by FuseBatchNormWithLinearPass (each applied once)."""
@@ -671,12 +662,6 @@ def test__no_pre_or_post_convert_passes(self):
         assert qr.pre_convert_passes is None
         assert qr.post_convert_passes is None
 
-    def test__post_prepare_passes_include_histogram_observer_fix(self):
-        """PTQ recipe post_prepare_passes contains _histogram_observer_fix_pass."""
-        qr = self._get_ptq_recipe().quantization_recipe
-        assert qr.post_prepare_passes is not None
-        assert _histogram_observer_fix_pass in qr.post_prepare_passes
-
     def test__post_prepare_passes_do_not_include_add_bn_fusion(self):
         """PTQ recipe post_prepare_passes must NOT contain AddSimulatedLinearBatchNormFusionQATPass."""
         qr = self._get_ptq_recipe().quantization_recipe
@@ -685,6 +670,13 @@ def test__post_prepare_passes_do_not_include_add_bn_fusion(self):
                 p, "__qualname__", ""
             )
 
+    def test__post_prepare_passes_is_empty(self):
+        """PTQ recipe post_prepare_passes is empty: NeutronQuantizer itself suppresses
+        quantization annotation for non-float nodes, so no post-prepare fixup pass is
+        needed (previously _histogram_observer_fix_pass was appended here)."""
+        qr = self._get_ptq_recipe().quantization_recipe
+        assert qr.post_prepare_passes == []
+
 
 # ---------------------------------------------------------------------------
 # QAT recipe: e2e test on a real model - equivalent to imperative QAT tests