NXP backend: Add pass for adding batch dimension for 3D input avgpool…

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)

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
 

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"

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"

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"

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"

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(),

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
 

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):

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()

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()