diff --git a/backends/nxp/aten_passes/add_batch_size_for_3d_input_pool_2d_ops.py b/backends/nxp/aten_passes/add_batch_size_for_3d_input_pool_2d_ops.py new file mode 100644 index 00000000000..7435b3b6969 --- /dev/null +++ b/backends/nxp/aten_passes/add_batch_size_for_3d_input_pool_2d_ops.py @@ -0,0 +1,116 @@ +# 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 executorch.backends.nxp.backend.edge_helper import input_rank + +from torch.fx import GraphModule, Node +from torch.fx.passes.infra.pass_base import PassBase, PassResult + + +class AddBatchSizeFor3DInputPool2DOps(PassBase): + """Adds batch size dimension for aten.adaptive_avg_pool2d.default, aten.avg_pool2d.default + and aten.max_pool2d.default ops with 3D input, as the Neutron Converter is unable to convert these ops with 3D input. + + │ + ┌──────▼──────┐ + │ reshape │ + │ (add batch) │ + └──────┬──────┘ + │ │ + ┌──────▼───────┐ ┌──────▼───────┐ + │ adaptive │ │ adaptive │ + │ _avg_pool2d/ │ replace with │ _avg_pool2d/ │ + │ avg_pool2d/ │ ──────────────► │ avg_pool2d/ │ + │ max_pool2d │ │ max_pool2d │ + │ 3D input │ │ 4D input │ + │ (C,H,W) │ │ (1,C H,W) │ + └──────┬───────┘ └──────┬───────┘ + │ │ + ┌──────▼───────┐ + │ reshape │ + │(remove batch)│ + └──────┬───────┘ + │ + """ + + module: GraphModule + + def _create_reshape_node(self, pool_node: Node, *args) -> Node: + reshape_node = self.module.graph.call_function( + torch.ops.aten.reshape.default, + args=args, + kwargs={}, + ) + reshape_node.meta["source_fn_stack"] = pool_node.meta.get("source_fn_stack", []) + input_node = args[0] + if input_node == pool_node: + # insert after pool_node + reshape_node.meta["val"] = input_node.meta["val"].squeeze(0) + else: + # insert before pool_node + reshape_node.meta["val"] = input_node.meta["val"].unsqueeze(0) + return reshape_node + + def call(self, module: GraphModule) -> PassResult: + self.module = module + + def _is_3d_pool(node_: Node) -> bool: + if node_.op != "call_function": + return False + + if node_.target not in [ + torch.ops.aten.adaptive_avg_pool2d.default, + torch.ops.aten.avg_pool2d.default, + torch.ops.aten.max_pool2d.default, + ]: + return False + + # Check if input is 3D (C, H, W) + rank = input_rank(node_, 0) + return rank == 3 + + made_changes = False + + for node in module.graph.nodes: + if not _is_3d_pool(node): + continue + + pool_node = node + input_node = pool_node.args[0] + + # Get input shape (C, H, W) + input_shape = input_node.meta["val"].shape + + # Get output shape (3D) before we modify metadata + output_shape_3d = pool_node.meta["val"].shape + + # Insert reshape to add batch dimension (1, C, H, W) + with module.graph.inserting_before(pool_node): + reshape_add_batch = self._create_reshape_node( + pool_node, input_node, [1, *input_shape] + ) + + # Update pool_node to use 4D input + pool_node.args = (reshape_add_batch, *pool_node.args[1:]) + + # Update pool_node output metadata to 4D + pool_node.meta["val"] = pool_node.meta["val"].unsqueeze(0) + + # Insert reshape to remove batch dimension AFTER pool_node + with module.graph.inserting_after(pool_node): + reshape_remove_batch = self._create_reshape_node( + pool_node, pool_node, list(output_shape_3d) + ) + + # Replace all uses of pool_node with reshape_remove_batch (except reshape_remove_batch itself) + pool_node.replace_all_uses_with(reshape_remove_batch) + # Restore the connection: reshape_remove_batch should use pool_node as input + reshape_remove_batch.update_arg(0, pool_node) + + made_changes = True + + return PassResult(module, made_changes) diff --git a/backends/nxp/aten_passes/decompose_split_to_slices_pass.py b/backends/nxp/aten_passes/decompose_split_to_slices_pass.py index d79397481b2..f975406a87a 100644 --- a/backends/nxp/aten_passes/decompose_split_to_slices_pass.py +++ b/backends/nxp/aten_passes/decompose_split_to_slices_pass.py @@ -3,7 +3,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from typing import Optional, TypeAlias +from typing import TypeAlias import torch from torch._subclasses import FakeTensor, FakeTensorMode @@ -150,7 +150,7 @@ def _replace_split_with_slices(self, input_node, split_node, starts, ends, dim): split_node.replace_all_uses_with(input_node) self.graph_module.graph.erase_node(split_node) - def call(self, graph_module: GraphModule) -> Optional[PassResult]: + def call(self, graph_module: GraphModule) -> PassResult: self.graph_module = graph_module made_changes = False diff --git a/backends/nxp/aten_passes/fuse_batch_norm_with_conv_pass.py b/backends/nxp/aten_passes/fuse_batch_norm_with_conv_pass.py index 33d78791b59..77d4ac0244a 100644 --- a/backends/nxp/aten_passes/fuse_batch_norm_with_conv_pass.py +++ b/backends/nxp/aten_passes/fuse_batch_norm_with_conv_pass.py @@ -2,7 +2,6 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from typing import Optional import torch from torch.export.unflatten import _assign_attr, _AttrKind @@ -54,7 +53,7 @@ def _get_tensor_constant_from_node(self, graph_module, node) -> Parameter | None attr_itr = getattr(attr_itr, atom) return attr_itr - def call(self, graph_module: GraphModule) -> Optional[PassResult]: + def call(self, graph_module: GraphModule) -> PassResult: def _is_batch_norm(node_: Node) -> bool: return ( node_.op == "call_function" diff --git a/backends/nxp/aten_passes/fuse_batch_norm_with_linear_pass.py b/backends/nxp/aten_passes/fuse_batch_norm_with_linear_pass.py index 6f95a17cc68..69502159b14 100644 --- a/backends/nxp/aten_passes/fuse_batch_norm_with_linear_pass.py +++ b/backends/nxp/aten_passes/fuse_batch_norm_with_linear_pass.py @@ -69,7 +69,7 @@ def _get_tensor_constant_from_node(self, graph_module, node) -> Parameter | None attr_itr = getattr(attr_itr, atom) return attr_itr - def call(self, graph_module: GraphModule) -> PassResult | None: + def call(self, graph_module: GraphModule) -> PassResult: def _is_batch_norm(node_: Node) -> bool: return ( node_.op == "call_function" diff --git a/backends/nxp/aten_passes/fuse_linear_and_add_pass.py b/backends/nxp/aten_passes/fuse_linear_and_add_pass.py index 20a32c1bcac..c74e8974dec 100644 --- a/backends/nxp/aten_passes/fuse_linear_and_add_pass.py +++ b/backends/nxp/aten_passes/fuse_linear_and_add_pass.py @@ -3,8 +3,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from typing import Optional - import torch from executorch.backends.nxp.backend.edge_helper import ( @@ -141,7 +139,7 @@ def _fuse_without_existing_bias( ) return True - def call(self, graph_module: GraphModule) -> Optional[PassResult]: + def call(self, graph_module: GraphModule) -> PassResult: def _is_applicable_linear_node(node_: Node): is_linear = ( node_.op == "call_function" diff --git a/backends/nxp/aten_passes/move_activation_before_concat.py b/backends/nxp/aten_passes/move_activation_before_concat.py index 8ba306d42e2..4ae288fbc1e 100644 --- a/backends/nxp/aten_passes/move_activation_before_concat.py +++ b/backends/nxp/aten_passes/move_activation_before_concat.py @@ -36,7 +36,7 @@ class MoveActivationBeforeConcat(PassBase): def __init__(self, neutron_target_spec: NeutronTargetSpec): self.neutron_target_spec = neutron_target_spec - def call(self, module: GraphModule) -> bool: + def call(self, module: GraphModule) -> PassResult: def _is_concat(node_: Node) -> bool: return ( node_.op == "call_function" diff --git a/backends/nxp/aten_passes/neutron_aten_pass_manager.py b/backends/nxp/aten_passes/neutron_aten_pass_manager.py index a20636ff9f2..8a7716f814c 100644 --- a/backends/nxp/aten_passes/neutron_aten_pass_manager.py +++ b/backends/nxp/aten_passes/neutron_aten_pass_manager.py @@ -7,6 +7,9 @@ import torch +from executorch.backends.nxp.aten_passes.add_batch_size_for_3d_input_pool_2d_ops import ( + AddBatchSizeFor3DInputPool2DOps, +) from executorch.backends.nxp.aten_passes.convert_1d_conv_to_2d import ( ConvertConv1dToConv2dPass, ) @@ -48,6 +51,7 @@ def _get_default_passes(neutron_target_spec, qat_mode: bool = False) -> list[PassType]: passes = [ + AddBatchSizeFor3DInputPool2DOps(), DecomposeSplitToSlicesPass(), SplitGroupConvolution(), SplitGRUBasedOnNumLayers(), diff --git a/backends/nxp/aten_passes/remove_nodes_with_known_outputs.py b/backends/nxp/aten_passes/remove_nodes_with_known_outputs.py index 3c08ac6c3fb..bcd85a9e415 100644 --- a/backends/nxp/aten_passes/remove_nodes_with_known_outputs.py +++ b/backends/nxp/aten_passes/remove_nodes_with_known_outputs.py @@ -150,7 +150,7 @@ def data_matches_meta_of_following_getitem_nodes( for get_item in users ) - def call(self, module: GraphModule) -> bool: + def call(self, module: GraphModule) -> PassResult: self.module = module made_changes = False diff --git a/backends/nxp/aten_passes/split_group_convolution.py b/backends/nxp/aten_passes/split_group_convolution.py index 22fc3a83cee..0a8fb07b3a4 100644 --- a/backends/nxp/aten_passes/split_group_convolution.py +++ b/backends/nxp/aten_passes/split_group_convolution.py @@ -193,7 +193,7 @@ def _create_parameter_node_for_data( return static_parameter_node - def call(self, module: GraphModule): + def call(self, module: GraphModule) -> PassResult: self.module = module def _is_conv(node_: Node): diff --git a/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py b/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py index 20e6814224e..1ed9a9f607c 100644 --- a/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py +++ b/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py @@ -6,9 +6,9 @@ from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.pass_base import PassResult from executorch.exir.passes import dead_code_elimination_pass from torch.fx import GraphModule +from torch.fx.passes.infra.pass_base import PassResult class RemoveUselessAsStridedCopyNodes(NeutronEdgePass): @@ -57,7 +57,7 @@ def _fold_as_strided_copy( return made_changes - def run(self, graph_module: GraphModule): + def run(self, graph_module: GraphModule) -> PassResult: made_changes = self._fold_as_strided_copy(graph_module) graph_module.recompile() diff --git a/backends/nxp/edge_passes/remove_io_quant_ops_pass.py b/backends/nxp/edge_passes/remove_io_quant_ops_pass.py index d49b646d489..a87eac7360c 100644 --- a/backends/nxp/edge_passes/remove_io_quant_ops_pass.py +++ b/backends/nxp/edge_passes/remove_io_quant_ops_pass.py @@ -69,7 +69,7 @@ def _get_quantizable_output_indices(self): return outputs_to_quantization - def call(self, graph_module: torch.fx.GraphModule): + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: input_indices = self._get_quantizable_input_indices() output_indices = self._get_quantizable_output_indices() diff --git a/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py b/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py new file mode 100644 index 00000000000..b426c34c260 --- /dev/null +++ b/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py @@ -0,0 +1,288 @@ +# 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.add_batch_size_for_3d_input_pool_2d_ops import ( + AddBatchSizeFor3DInputPool2DOps, +) +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.model_output_comparator import ( + AllCloseOutputComparator, +) +from executorch.backends.nxp.tests.models import ( + AdaptiveAvgPool2dModule, + AvgPool2dModule, + MaxPool2dModule, +) +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.ops_aliases import ( + AdaptiveAvgPool2D, + AvgPool2D, + GetItem, + MaxPool2DWithIndices, + ViewCopy, +) + + +@pytest.fixture(autouse=True) +def reseed_model_per_test_run(): + torch.manual_seed(42) + np.random.seed(23) + + +def apply_individual_pass_and_compare(model, export_input, pool_op): + exir_program_aten = torch.export.export( + model, + export_input, + ).module() + + # Check that pool op is present and has 3D input + assert graph_contains_any_of_ops(exir_program_aten.graph, [pool_op]) + nodes = list(exir_program_aten.graph.nodes) + pool_node = nodes[1] + assert pool_node.target == pool_op + assert len(pool_node.meta["val"].shape) == 3 + pool_count_prev = sum( + [n.target == pool_op for n in list(exir_program_aten.graph.nodes)] + ) + + # Check that reshape is not present before the pass + reshape_count_before = sum( + [ + n.target == torch.ops.aten.reshape.default + for n in list(exir_program_aten.graph.nodes) + ] + ) + assert reshape_count_before == 0 + + outputs_before = [o.detach().numpy() for o in exir_program_aten(*export_input)] + + # Apply the optimization. + NeutronAtenPassManager(neutron_target_spec, [AddBatchSizeFor3DInputPool2DOps()])( + exir_program_aten + ) + + # Make sure pool op is still in the model and has 4D input with batch size == 1 + assert graph_contains_any_of_ops(exir_program_aten.graph, [pool_op]) + nodes = list(exir_program_aten.graph.nodes) + pool_node = nodes[2] + assert pool_node.target == pool_op + assert len(pool_node.meta["val"].shape) == 4 and pool_node.meta["val"].shape[0] == 1 + + # Make sure there is `reshape` in the model. + assert graph_contains_any_of_ops( + exir_program_aten.graph, + [torch.ops.aten.reshape.default], + ) + + reshape_count_after = sum( + [ + n.target == torch.ops.aten.reshape.default + for n in list(exir_program_aten.graph.nodes) + ] + ) + + pool_count_after = sum( + [n.target == pool_op for n in list(exir_program_aten.graph.nodes)] + ) + + # Make sure the number of pool operators is the same + assert pool_count_prev == pool_count_after + + # Make sure we added 2 reshape operations per pool operation (add batch + remove batch) + assert reshape_count_after == 2 * pool_count_after + + outputs_after = [o.detach().numpy() for o in exir_program_aten(*export_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], rtol=1e-5, atol=1e-5) + + +class TestAddBatchSizeFor3DPoolOps: + @pytest.mark.parametrize( + "input_shape, output_size", + [ + pytest.param((16, 32, 32), (16, 16), id="3D, output_size=(16, 16)"), + pytest.param((32, 64, 64), (8, 8), id="3D, output_size=(8, 8)"), + ], + ) + def test_add_batch_size_adaptive_avgpool2d(self, input_shape, output_size): + model = AdaptiveAvgPool2dModule(output_size=output_size) + example_input = torch.rand(input_shape, dtype=torch.float32) + apply_individual_pass_and_compare( + model, (example_input,), torch.ops.aten.adaptive_avg_pool2d.default + ) + + @pytest.mark.parametrize( + "input_shape, kernel_size, stride, padding", + [ + pytest.param( + (16, 32, 32), 2, 2, 0, id="3D, kernel=2, stride=2, no padding" + ), + pytest.param((32, 64, 64), 3, 3, 1, id="3D, kernel=3, stride=3, padding=1"), + pytest.param( + (8, 16, 16), (2, 2), (2, 2), 0, id="3D, kernel=(2,2), stride=(2,2)" + ), + ], + ) + def test_add_batch_size_avgpool2d(self, input_shape, kernel_size, stride, padding): + model = AvgPool2dModule(kernel_size=kernel_size, stride=stride, padding=padding) + example_input = torch.rand(input_shape, dtype=torch.float32) + apply_individual_pass_and_compare( + model, (example_input,), torch.ops.aten.avg_pool2d.default + ) + + @pytest.mark.parametrize( + "input_shape, kernel_size, stride, padding", + [ + pytest.param( + (16, 32, 32), 2, 2, 0, id="3D, kernel=2, stride=2, no padding" + ), + pytest.param((32, 64, 64), 3, 3, 1, id="3D, kernel=3, stride=3, padding=1"), + pytest.param( + (8, 16, 16), (2, 2), (2, 2), 0, id="3D, kernel=(2,2), stride=(2,2)" + ), + ], + ) + def test_add_batch_size_maxpool2d(self, input_shape, kernel_size, stride, padding): + model = MaxPool2dModule(kernel_size=kernel_size, stride=stride, padding=padding) + example_input = torch.rand(input_shape, dtype=torch.float32) + apply_individual_pass_and_compare( + model, (example_input,), torch.ops.aten.max_pool2d.default + ) + + @pytest.mark.parametrize( + "input_shape", + [ + pytest.param((1, 16, 32, 32), id="4D input - should not transform"), + pytest.param( + (2, 16, 32, 32), id="4D input with batch=2 - should not transform" + ), + ], + ) + @pytest.mark.parametrize( + "model, pool_op", + [ + pytest.param( + AdaptiveAvgPool2dModule(output_size=(16, 16)), + torch.ops.aten.adaptive_avg_pool2d.default, + id="AdaptiveAvgPoolModel", + ), + pytest.param( + AvgPool2dModule(kernel_size=2, stride=2, padding=0), + torch.ops.aten.avg_pool2d.default, + id="AvgPoolModel", + ), + pytest.param( + MaxPool2dModule(kernel_size=2, stride=2, padding=0), + torch.ops.aten.max_pool2d.default, + id="MaxPoolModel", + ), + ], + ) + def test_no_transform_for_4d_input(self, input_shape, model, pool_op): + example_input = torch.rand(input_shape, dtype=torch.float32) + + exir_program_aten = torch.export.export( + model, + (example_input,), + ).module() + + # Check that pool op is present + assert graph_contains_any_of_ops( + exir_program_aten.graph, + [pool_op], + ) + + # Check that reshape is not present before the pass + reshape_count_before = sum( + [ + n.target == torch.ops.aten.reshape.default + for n in list(exir_program_aten.graph.nodes) + ] + ) + assert reshape_count_before == 0 + + # Apply the optimization. + NeutronAtenPassManager( + neutron_target_spec, [AddBatchSizeFor3DInputPool2DOps()] + )(exir_program_aten) + + # Check that reshape count hasn't changed (no transformation for 4D input) + reshape_count_after = sum( + [ + n.target == torch.ops.aten.reshape.default + for n in list(exir_program_aten.graph.nodes) + ] + ) + + assert reshape_count_before == reshape_count_after + + @pytest.mark.parametrize( + "input_shape", + [ + (16, 32, 32), + (32, 64, 64), + ], + ids=lambda shape: f"3D_{shape[0]}x{shape[1]}x{shape[2]}", + ) + @pytest.mark.parametrize( + "pool_type", + ["adaptive_avg", "avg", "max"], + ids=lambda pool_type: f"{pool_type}pool2d", + ) + def test__3d_pool__full_pipeline( + self, mocker, request, input_shape: tuple[int, ...], pool_type: str + ): + expected_delegated_ops = {} + match pool_type: + case "adaptive_avg": + model = AdaptiveAvgPool2dModule(output_size=(16, 16)) + expected_delegated_ops = {ViewCopy: 2, AdaptiveAvgPool2D: 1} + case "avg": + model = AvgPool2dModule(kernel_size=2, stride=2, padding=0) + expected_delegated_ops = {ViewCopy: 2, AvgPool2D: 1} + case _: + model = MaxPool2dModule(kernel_size=2, stride=2, padding=0) + expected_delegated_ops = { + ViewCopy: 2, + MaxPool2DWithIndices: 1, + GetItem: 1, + } + + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops=expected_delegated_ops, + expected_non_delegated_ops={}, + ) + + dataset_creator = RandomDatasetCreator(low=-1, high=1) + + remove_quant_io_ops = True # Use quantized dataset. + output_comparator = AllCloseOutputComparator(atol=1) # Allow single bit error. + + lower_run_compare( + model, + input_shape, + graph_verifier, + request, + dataset_creator, + output_comparator, + remove_quant_io_ops=remove_quant_io_ops, + ) diff --git a/backends/nxp/tests/models.py b/backends/nxp/tests/models.py index d92002f47a7..f075157b7cb 100644 --- a/backends/nxp/tests/models.py +++ b/backends/nxp/tests/models.py @@ -400,11 +400,11 @@ def forward(self, x): class MaxPool2dModule(torch.nn.Module): - def __init__(self, padding=0): + def __init__(self, padding=0, kernel_size=3, stride=2): super().__init__() self.max_pool2d = torch.nn.MaxPool2d( - kernel_size=3, stride=2, padding=padding, dilation=1 + kernel_size=kernel_size, stride=stride, padding=padding, dilation=1 ) def forward(self, x): @@ -426,7 +426,7 @@ def forward(self, x): class AvgPool2dModule(torch.nn.Module): - def __init__(self, count_include_pad, padding=0, kernel_size=3, stride=2): + def __init__(self, count_include_pad=True, padding=0, kernel_size=3, stride=2): super().__init__() self.avg_pool = torch.nn.AvgPool2d(