Arm backend: Support rank-3 max_pool2d inputs (pytorch#22441)

Normalize unbatched max_pool2d inputs from [C, H, W] to [1, C, H, W]
before decomposition, then restore the rank-3 output.

Enable rank-3 support checks and add FP, INT, and VGF test coverage,
including quantized VGF execution.

Change-Id: Ib8852053b5248601d59685669523079fa19c0955
Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>

Author: YufengShi-dudu

backends/arm/_passes/__init__.py

@@ -170,6 +170,9 @@
 from .normalize_index_put_none_indices_pass import (  # noqa
     NormalizeIndexPutNoneIndicesPass,
 )
+from .normalize_max_pool2d_input_rank_pass import (  # noqa
+    NormalizeMaxPool2dInputRankPass,
+)
 from .normalize_while_initial_args_pass import NormalizeWhileInitialArgsPass  # noqa
 from .promote_bool_operands_pass import PromoteBoolOperandsPass  # noqa
 from .propagate_view_copy_permute_pass import (  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -138,6 +138,7 @@
     NormalizeDelegateIOLayoutPass,
     NormalizeIndexPutBoolIndexTensorPass,
     NormalizeIndexPutNoneIndicesPass,
+    NormalizeMaxPool2dInputRankPass,
     NormalizeTransformInputPlaceholdersPass,
     NormalizeWhileInitialArgsPass,
     PromoteBoolOperandsPass,
@@ -637,6 +638,7 @@ def _tosa_pipeline(
                 UnsqueezeBeforeRepeatPass(),
                 DecomposeCumsumPass(exported_program),
                 DecomposeAsStridedCopyPass(),
+                NormalizeMaxPool2dInputRankPass(),
                 DecomposeMaxPool2dPass(),
                 DecomposeLargeStrideMaxPool2dForU55Pass(),
                 SizeAdjustInputPass(),

backends/arm/_passes/normalize_max_pool2d_input_rank_pass.py

@@ -0,0 +1,97 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# 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 Set, Type
+
+import torch
+from executorch.backends.arm._passes import ArmOpTargetedPass
+from executorch.backends.arm._passes.arm_pass_utils import (
+    create_node,
+    get_first_fake_tensor,
+)
+from executorch.backends.arm._passes.convert_squeezes_to_view import (
+    ConvertSqueezesToViewPass,
+)
+from executorch.backends.arm._passes.decompose_maxpool2d_with_dilation_pass import (
+    DecomposeMaxPool2dPass,
+)
+from executorch.backends.arm._passes.rewrite_max_pool2d_pass import RewriteMaxPool2dPass
+from executorch.backends.arm._passes.size_adjust_input_pass import SizeAdjustInputPass
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, PassResult
+from torch.fx import Node
+
+
+class NormalizeMaxPool2dInputRankPass(ArmOpTargetedPass):
+    """Normalize unbatched rank-3 max_pool2d inputs to rank 4.
+
+    Unsqueeze inputs from ``[C, H, W]`` to ``[1, C, H, W]``.
+
+    Squeeze the leading dimension after pooling to restore the rank-3 output.
+
+    The complete shape transformation is::
+
+        [C, H, W]
+          -> unsqueeze(0) -> [1, C, H, W]
+          -> max_pool2d -> [1, C, H_out, W_out]
+          -> squeeze(0) -> [C, H_out, W_out]
+
+    """
+
+    target_ops = (exir_ops.edge.aten.max_pool2d.default,)
+    _passes_required_after: Set[Type[ExportPass]] = {
+        ConvertSqueezesToViewPass,
+        DecomposeMaxPool2dPass,
+        RewriteMaxPool2dPass,
+        SizeAdjustInputPass,
+    }
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        graph = graph_module.graph
+        modified = False
+        pool_nodes = graph.find_nodes(op="call_function", target=self.target_ops[0])
+
+        for pool_node in pool_nodes:
+            input_node = pool_node.args[0]
+            if not isinstance(input_node, Node):
+                raise RuntimeError("Expected max_pool2d input to be a node")
+
+            input_fake = get_first_fake_tensor(input_node)
+            if input_fake.dim() != 3:
+                continue
+
+            output_fake = get_first_fake_tensor(pool_node)
+            with graph.inserting_before(pool_node):
+                unsqueeze = create_node(
+                    graph,
+                    exir_ops.edge.aten.unsqueeze_copy.default,
+                    args=(input_node, 0),
+                    from_node=pool_node,
+                    inherit_qparams=False,
+                )
+                unsqueeze.meta["val"] = input_fake.unsqueeze(0)
+                pool_node.replace_input_with(input_node, unsqueeze)
+
+            pool_node.meta["val"] = output_fake.unsqueeze(0)
+            original_users = list(pool_node.users)
+            with graph.inserting_after(pool_node):
+                squeeze = create_node(
+                    graph,
+                    exir_ops.edge.aten.squeeze_copy.dims,
+                    args=(pool_node, [0]),
+                    from_node=pool_node,
+                    inherit_qparams=False,
+                )
+                squeeze.meta["val"] = output_fake
+                for user in original_users:
+                    user.replace_input_with(pool_node, squeeze)
+
+            modified = True
+
+        if modified:
+            graph.lint()
+            graph_module.recompile()
+
+        return PassResult(graph_module, modified)

backends/arm/operator_support/pool_2d_support.py

@@ -228,6 +228,14 @@ def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
         """
 
         shape = cast(torch.Tensor, node.all_input_nodes[0].meta["val"]).shape
+        if len(shape) == 3:
+            shape = torch.Size((1, *shape))
+        elif len(shape) != 4:
+            self.reporter.report_reject(
+                node, f"Maxpool2d needs rank 3 or 4 input, got shape {list(shape)}"
+            )
+            return False
+
         kernel = cast(tuple[int, int], node.args[1])
         stride = cast(tuple[int, int], node.args[2])
         padding = cast(tuple[int, int], node.args[3]) if len(node.args) >= 4 else (0, 0)

backends/arm/test/ops/test_max_pool.py

@@ -82,6 +82,12 @@
         [3, 2, 1],
     ),
 }
+
+test_data_suite_rank3 = {
+    "bev_class_slice": lambda: (torch.rand(1, 8, 8), [1, 1, 0]),
+    "spatial_pool": lambda: (torch.rand(3, 9, 11), [3, 2, 1]),
+}
+
 test_data_suite_fp8 = {
     "rand_fp8e4m3": lambda: (
         torch.rand(1, 8, 20, 20).to(torch.float8_e4m3fn),
@@ -168,7 +174,11 @@ def forward(self, x):
 
 
 @common.parametrize(
-    "test_data", test_data_suite | test_data_suite_fp16 | test_data_suite_bf16
+    "test_data",
+    test_data_suite
+    | test_data_suite_rank3
+    | test_data_suite_fp16
+    | test_data_suite_bf16,
 )
 def test_max_pool2d_tosa_FP(test_data: torch.Tensor):
     test_data, model_params = test_data()
@@ -197,7 +207,7 @@ def test_max_pool2d_tosa_FP_fp8(test_data: torch.Tensor):
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data_suite)
+@common.parametrize("test_data", test_data_suite | test_data_suite_rank3)
 def test_max_pool2d_tosa_INT(test_data: torch.Tensor):
     test_data, model_params = test_data()
     pipeline = TosaPipelineINT[input_t1](
@@ -374,22 +384,28 @@ def test_max_pool2d_tosa_INT_dilation(test_data):
 
 # VGF tests
 @common.parametrize(
-    "test_data", test_data_suite | test_data_suite_bf16 | test_data_suite_fp16
+    "test_data",
+    test_data_suite
+    | test_data_suite_rank3
+    | test_data_suite_bf16
+    | test_data_suite_fp16,
 )
 @common.SkipIfNoModelConverter
 def test_max_pool2d_vgf_no_quant(test_data: torch.Tensor):
     test_data, model_params = test_data()
+    run_on_vulkan_runtime = test_data.dim() == 4
     pipeline = VgfPipeline[input_t1](
         MaxPool2d(*model_params),
         (test_data,),
         aten_op,
         exir_op,
         quantize=False,
+        run_on_vulkan_runtime=run_on_vulkan_runtime,
     )
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data_suite)
+@common.parametrize("test_data", test_data_suite | test_data_suite_rank3)
 @common.SkipIfNoModelConverter
 def test_max_pool2d_vgf_quant(test_data: torch.Tensor):
     test_data, model_params = test_data()

backends/arm/test/passes/test_normalize_max_pool2d_input_rank_pass.py

@@ -0,0 +1,87 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# 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 Tuple
+
+import torch
+from executorch.backends.arm._passes import (
+    NormalizeMaxPool2dInputRankPass,
+    RemoveGetItemPass,
+)
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+from executorch.backends.test.harness.stages import StageType
+from executorch.exir.dialects._ops import ops as exir_ops
+
+
+input_t = Tuple[torch.Tensor]
+
+
+class MaxPool2d(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.max_pool2d(
+            x,
+            kernel_size=(3, 2),
+            stride=(2, 1),
+            padding=(1, 0),
+            dilation=(1, 1),
+            ceil_mode=True,
+        )
+
+
+def test_normalize_rank3_max_pool2d_input() -> None:
+    pipeline = PassPipeline[input_t](
+        MaxPool2d(),
+        (torch.rand(3, 9, 11),),
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_max_pool2d_with_indices_default": 1,
+        },
+        ops_after_pass={
+            "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default": 1,
+            "executorch_exir_dialects_edge__ops_aten_max_pool2d_default": 1,
+            "executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims": 1,
+        },
+        pass_list=[RemoveGetItemPass, NormalizeMaxPool2dInputRankPass],
+    )
+    pipeline.run()
+
+    exported_program = pipeline.tester.get_artifact(
+        StageType.RUN_PASSES
+    ).exported_program()
+    pool_node = next(
+        node
+        for node in exported_program.graph.nodes
+        if node.target == exir_ops.edge.aten.max_pool2d.default
+    )
+    unsqueeze_node = pool_node.args[0]
+    assert isinstance(unsqueeze_node, torch.fx.Node)
+    assert unsqueeze_node.target == exir_ops.edge.aten.unsqueeze_copy.default
+    assert unsqueeze_node.args[1] == 0
+    assert tuple(pool_node.args[1]) == (3, 2)
+    assert tuple(pool_node.args[2]) == (2, 1)
+    assert tuple(pool_node.args[3]) == (1, 0)
+    assert tuple(pool_node.args[4]) == (1, 1)
+    assert pool_node.args[5] is True
+
+    squeeze_node = next(iter(pool_node.users))
+    assert squeeze_node.target == exir_ops.edge.aten.squeeze_copy.dims
+    assert squeeze_node.args == (pool_node, [0])
+
+
+def test_normalize_rank4_max_pool2d_input_is_noop() -> None:
+    PassPipeline[input_t](
+        MaxPool2d(),
+        (torch.rand(1, 3, 9, 11),),
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_max_pool2d_with_indices_default": 1,
+        },
+        ops_after_pass={
+            "executorch_exir_dialects_edge__ops_aten_max_pool2d_default": 1,
+        },
+        ops_not_after_pass=[
+            "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default",
+            "executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims",
+        ],
+        pass_list=[RemoveGetItemPass, NormalizeMaxPool2dInputRankPass],
+    ).run()