diff --git a/backends/nxp/aten_passes/convert_scalar_to_attr.py b/backends/nxp/aten_passes/convert_scalar_to_attr.py new file mode 100644 index 00000000000..4fd75485109 --- /dev/null +++ b/backends/nxp/aten_passes/convert_scalar_to_attr.py @@ -0,0 +1,100 @@ +# 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 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. + """ + + @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 torch.float32 + + def _create_scalar_attr_node(self, node: Node, scalar_value: int | float) -> Node: + # 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 + + 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_args[i] = self._create_scalar_attr_node(node, 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) diff --git a/backends/nxp/aten_passes/neutron_aten_pass_manager.py b/backends/nxp/aten_passes/neutron_aten_pass_manager.py index d00b45be582..a20636ff9f2 100644 --- a/backends/nxp/aten_passes/neutron_aten_pass_manager.py +++ b/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: diff --git a/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py b/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py new file mode 100644 index 00000000000..ddecfb9bb81 --- /dev/null +++ b/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py @@ -0,0 +1,112 @@ +# 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.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", + [ + pytest.param(AddScalarModule, torch.ops.aten.add.Tensor, id="Add op."), + pytest.param(SubScalarModule, torch.ops.aten.sub.Tensor, id="Sub op."), + pytest.param(MulScalarModule, torch.ops.aten.mul.Tensor, id="Mul op."), + ], + ) + def test__scalar_to_attr(self, model_cls, expected_op): + input_shape = (2, 5, 7, 9) + model = model_cls(scalar=2.0) + + 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 + ) + + 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, + ) + diff --git a/backends/nxp/tests/models.py b/backends/nxp/tests/models.py index a00ef602395..04e77f2d282 100644 --- a/backends/nxp/tests/models.py +++ b/backends/nxp/tests/models.py @@ -989,6 +989,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__()