NXP backend: added support for scalar args of elementary ops (pytorch#21011)

### Summary

Added support for scalar args of elementary ops.

### Test plan

tests can be manually run using `pytest -c /dev/null
backends/nxp/tests/`

cc @robert-kalmar @JakeStevens @digantdesai @rascani @MartinPavella

Author: novak-vaclav

backends/nxp/aten_passes/convert_scalar_to_attr.py

@@ -0,0 +1,119 @@
+# 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
+
+import torch
+
+from torch.ao.quantization.fx.utils import get_new_attr_name_with_prefix
+
+from torch.export.unflatten import _assign_attr, _AttrKind
+from torch.fx import GraphModule, Node
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+
+
+class ConvertScalarToAttrPass(PassBase):
+    """
+    Convert scalar (Python `int`/`float`) arguments of elementwise operators into
+    `get_attr` nodes holding a constant tensor. For example `aten.mul.Tensor(x, 2.0)`
+    is rewritten so the scalar `2.0` becomes a `get_attr` node referencing a tensor
+    constant with the value `2.0`.
+
+                       x                                                     x     get_attr(2.0)
+                       |                                                     |    /
+        ┌──────────────▼─────────────┐     replace with     ┌────────────────▼───────────────┐
+        |   aten.mul.Tensor(x, 2.0)  |  ─────────────────►  |  aten.mul.Tensor(x, get_attr)  |
+        └──────────────┬─────────────┘                      └────────────────┬───────────────┘
+                       |                                                     |
+                       v                                                     v
+                      out                                                   out
+
+    If not done, the scalar arguments prevent from proper QDQ pattern utilization and
+    the operator cannot be delegated to Neutron.
+    """
+
+    _target_ops = (
+        torch.ops.aten.add.Tensor,
+        torch.ops.aten.sub.Tensor,
+        torch.ops.aten.mul.Tensor,
+    )
+
+    @staticmethod
+    def _get_ref_dtype(node: Node) -> torch.dtype:
+        # Infer the constant dtype from a sibling tensor input of the operator so
+        # the created constant matches the operator's tensor operand.
+        # Fall back to the node's own output dtype, then `float32`.
+        for arg in node.args:
+            if isinstance(arg, Node):
+                arg_val = arg.meta.get("val")
+                if isinstance(arg_val, torch.Tensor):
+                    return arg_val.dtype
+
+        node_val = node.meta.get("val")
+        if isinstance(node_val, torch.Tensor):
+            return node_val.dtype
+
+        return None
+
+    def _create_scalar_attr_node(self, node: Node, scalar_value: float) -> Node:
+        tensor_dtype = self._get_ref_dtype(node)
+        # Don't change anything if type cannot be inferred.
+        if tensor_dtype is None:
+            logging.warning(f"Cannot infer dtype of node `{node}`")
+            return None
+
+        # Create a constant tensor holding the scalar value.
+        tensor = torch.tensor(scalar_value, dtype=self._get_ref_dtype(node))
+        tensor_name = get_new_attr_name_with_prefix("_scalar_const_")(self.graph_module)
+
+        _assign_attr(
+            torch.nn.Parameter(tensor, requires_grad=False),
+            self.graph_module,
+            tensor_name,
+            _AttrKind.PARAMETER,
+        )
+
+        fake_mode = node.meta["val"].fake_mode
+        with self.graph_module.graph.inserting_before(node):
+            get_attr_node = self.graph_module.graph.create_node(
+                "get_attr", tensor_name, (), {}
+            )
+            get_attr_node.meta["val"] = fake_mode.from_tensor(
+                tensor, static_shapes=True
+            )
+
+        return get_attr_node
+
+    def call(self, graph_module: GraphModule) -> PassResult:
+        self.graph_module = graph_module
+        made_changes = False
+
+        for node in list(graph_module.graph.nodes):
+            if node.op != "call_function":
+                continue
+
+            if node.target not in self._target_ops:
+                continue
+
+            new_args = list(node.args)
+            node_changed = False
+            for i, arg in enumerate(new_args):
+                # Note: `bool` is subclass on `int`.
+                if isinstance(arg, bool) or not isinstance(arg, (int, float)):
+                    continue
+
+                new_arg = self._create_scalar_attr_node(node, arg)
+                if new_arg is not None:
+                    new_args[i] = new_arg
+                    node_changed = True
+
+            if node_changed:
+                node.args = tuple(new_args)
+                made_changes = True
+
+        self.graph_module.graph.eliminate_dead_code()
+        self.graph_module.recompile()
+
+        return PassResult(graph_module, made_changes)

backends/nxp/aten_passes/neutron_aten_pass_manager.py

@@ -11,6 +11,9 @@
     ConvertConv1dToConv2dPass,
 )
 from executorch.backends.nxp.aten_passes.convert_div_to_mul import ConvertDivToMulPass
+from executorch.backends.nxp.aten_passes.convert_scalar_to_attr import (
+    ConvertScalarToAttrPass,
+)
 from executorch.backends.nxp.aten_passes.decompose_split_to_slices_pass import (
     DecomposeSplitToSlicesPass,
 )
@@ -53,6 +56,7 @@ def _get_default_passes(neutron_target_spec, qat_mode: bool = False) -> list[Pas
         MoveActivationBeforeConcat(neutron_target_spec),
         ConvertDivToMulPass(),
         ConvertConv1dToConv2dPass(neutron_target_spec),
+        ConvertScalarToAttrPass(),
     ]
 
     if not qat_mode:

backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py

@@ -0,0 +1,129 @@
+# 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 numpy as np
+
+# noinspection PyUnusedImports
+import pytest
+import torch
+from executorch.backends.nxp.aten_passes.convert_scalar_to_attr import (
+    ConvertScalarToAttrPass,
+)
+from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import (
+    NeutronAtenPassManager,
+)
+
+from executorch.backends.nxp.backend.edge_helper import (
+    try_get_tensor_constant_from_node,
+)
+from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator
+from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec
+from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops
+from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier
+from executorch.backends.nxp.tests.models import (
+    AddScalarModule,
+    MulScalarModule,
+    SubScalarModule,
+)
+from executorch.backends.nxp.tests.nsys_testing import (
+    AllCloseOutputComparator,
+    lower_run_compare,
+)
+from executorch.backends.nxp.tests.ops_aliases import AddTensor, MulTensor, SubTensor
+
+
+@pytest.fixture(autouse=True)
+def reseed_model_per_test_run():
+    torch.manual_seed(42)
+    np.random.seed(23)
+
+
+class TestConvertScalarToAttr:
+    @pytest.mark.parametrize(
+        "model_cls, expected_op, scalar_val",
+        [
+            pytest.param(AddScalarModule, torch.ops.aten.add.Tensor, 2.0, id="Add op."),
+            pytest.param(SubScalarModule, torch.ops.aten.sub.Tensor, 2.0, id="Sub op."),
+            pytest.param(MulScalarModule, torch.ops.aten.mul.Tensor, 2.0, id="Mul op."),
+        ],
+    )
+    def test__scalar_to_attr(self, model_cls, expected_op, scalar_val):
+        input_shape = (2, 5, 7, 9)
+        model = model_cls(scalar=scalar_val)
+
+        example_input = torch.rand(input_shape)
+        exir_program_aten = torch.export.export(model, (example_input,)).module()
+
+        # Check if the node with scalar arg did not disappear.
+        assert graph_contains_any_of_ops(exir_program_aten.graph, [expected_op])
+        outputs_before = [o.detach().numpy() for o in exir_program_aten(example_input)]
+
+        # Apply the optimization.
+        NeutronAtenPassManager(neutron_target_spec, [ConvertScalarToAttrPass()])(
+            exir_program_aten
+        )
+
+        exp_op_node = [
+            n for n in exir_program_aten.graph.nodes if n.target == expected_op
+        ][0]
+        # Check that no arg is `float` or `int`.
+        # Note: `bool` is subtype of `int`, but `bool` does not need to be converted to `get_attr`.
+        assert not any(
+            isinstance(arg, (int, float)) and not isinstance(arg, bool)
+            for arg in exp_op_node.args
+        )
+        # Check that the scalar value is present.
+        assert any(
+            arg.op == "get_attr"
+            and try_get_tensor_constant_from_node(exir_program_aten, arg).item()
+            == scalar_val
+            for arg in exp_op_node.args
+        )
+
+        outputs_after = [o.detach().numpy() for o in exir_program_aten(example_input)]
+
+        # Make sure the model still produces the exact same output.
+        assert len(outputs_before) == len(outputs_after)
+        for i in range(len(outputs_before)):
+            assert np.allclose(outputs_before[i], outputs_after[i])
+
+    @pytest.mark.parametrize(
+        "model_cls, expected_op",
+        [
+            pytest.param(AddScalarModule, AddTensor, id="Add op."),
+            pytest.param(SubScalarModule, SubTensor, id="Sub op."),
+            pytest.param(MulScalarModule, MulTensor, id="Mul op."),
+        ],
+    )
+    def test__scalar_to_attr__full_pipeline(
+        self,
+        mocker,
+        request,
+        model_cls,
+        expected_op,
+    ):
+        input_shape = (2, 7, 5, 11)
+        model = model_cls(scalar=2.0)
+
+        graph_verifier = DetailedGraphVerifier(
+            mocker,
+            expected_delegated_ops={expected_op: 1},
+            expected_non_delegated_ops={},
+        )
+        dataset_creator = RandomDatasetCreator(low=-1.0, high=1.0)
+
+        # Quantize the dataset and allow a single bit error.
+        remove_quant_io_ops = True
+        comparator = AllCloseOutputComparator(atol=1)
+
+        lower_run_compare(
+            model,
+            input_shape,
+            graph_verifier,
+            request,
+            dataset_creator,
+            comparator,
+            remove_quant_io_ops=remove_quant_io_ops,
+        )

backends/nxp/tests/models.py

@@ -1020,6 +1020,33 @@ def forward(self, x, divisor):
         return x / divisor
 
 
+class AddScalarModule(torch.nn.Module):
+    def __init__(self, scalar: float | int = 2.0):
+        super().__init__()
+        self.scalar = scalar
+
+    def forward(self, x):
+        return x + self.scalar
+
+
+class MulScalarModule(torch.nn.Module):
+    def __init__(self, scalar: float | int = 2.0):
+        super().__init__()
+        self.scalar = scalar
+
+    def forward(self, x):
+        return x * self.scalar
+
+
+class SubScalarModule(torch.nn.Module):
+    def __init__(self, scalar: float | int = 2.0):
+        super().__init__()
+        self.scalar = scalar
+
+    def forward(self, x):
+        return x - self.scalar
+
+
 class BatchMatMulModel(torch.nn.Module):
     def __init__(self):
         super().__init__()