Legalize large-stride non-overlapping maxpool on U55

backends/arm/_passes/__init__.py

@@ -1,3 +1,3 @@
 # Copyright 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -74,6 +74,9 @@
 from .decompose_linear_pass import DecomposeLinearPass  # noqa
 from .decompose_log1p_pass import DecomposeLog1pPass  # noqa
 from .decompose_logit_pass import DecomposeLogitPass  # noqa
+from .decompose_large_stride_maxpool2d_pass import (  # noqa
+    DecomposeLargeStrideMaxPool2dPass,
+)
 from .decompose_lstm_pass import DecomposeLstmPass  # noqa
 from .decompose_masked_fill_pass import DecomposeMaskedFillPass  # noqa
 from .decompose_matmul import DecomposeMatmulPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 # Copyright 2024-2026 Arm Limited and/or its affiliates.
@@ -69,6 +69,7 @@
     DecomposeIndexTensorToGatherPass,
     DecomposeIntPowPass,
     DecomposeLayerNormPass,
+    DecomposeLargeStrideMaxPool2dPass,
     DecomposeLeakyReLUPass,
     DecomposeLinalgVectorNormPass,
     DecomposeLinearPass,
@@ -611,6 +612,7 @@
                 DecomposeCumsumPass(exported_program),
                 DecomposeAsStridedCopyPass(),
                 DecomposeMaxPool2dPass(),
+                DecomposeLargeStrideMaxPool2dPass(),
                 SizeAdjustInputPass(),
                 DecomposeUnsupportedBilinearResizePass(self.tosa_spec),
                 RewriteAdaptiveAvgPool2dPass(),

backends/arm/_passes/decompose_large_stride_maxpool2d_pass.py

@@ -0,0 +1,152 @@
+# 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.size_adjust_input_pass import SizeAdjustInputPass
+from executorch.backends.arm.tosa.specification import get_context_spec
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+
+_U55_MAX_POOL_STRIDE = 3
+
+
+def can_decompose_large_stride_maxpool2d(
+    kernel,
+    stride,
+    padding,
+    dilation,
+    ceil_mode,
+    input_shape,
+) -> bool:
+    kernel_h, kernel_w = (kernel, kernel) if isinstance(kernel, int) else kernel
+    stride_h, stride_w = (stride, stride) if isinstance(stride, int) else stride
+    padding_h, padding_w = (
+        (padding, padding) if isinstance(padding, int) else padding
+    )
+    dilation_h, dilation_w = (
+        (dilation, dilation) if isinstance(dilation, int) else dilation
+    )
+    height, width = input_shape[-2:]
+
+    return (
+        not isinstance(height, torch.SymInt)
+        and not isinstance(width, torch.SymInt)
+        and max(stride_h, stride_w) > _U55_MAX_POOL_STRIDE
+        and (kernel_h, kernel_w) == (stride_h, stride_w)
+        and (padding_h, padding_w) == (0, 0)
+        and (dilation_h, dilation_w) == (1, 1)
+        and not ceil_mode
+        and 1 <= kernel_h <= 256
+        and 1 <= kernel_w <= 256
+        and height >= kernel_h
+        and width >= kernel_w
+    )
+
+
+class DecomposeLargeStrideMaxPool2dPass(ArmOpTargetedPass):
+    """Legalize non-overlapping max_pool2d with strides unsupported by U55."""
+
+    _passes_required_after: Set[Type[ExportPass]] = {SizeAdjustInputPass}
+    target_ops = (exir_ops.edge.aten.max_pool2d.default,)
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in self.target_ops or not get_context_spec().is_U55_subset:
+            return super().call_operator(op, args, kwargs, meta)
+
+        x = args[0]
+        kernel = args[1]
+        stride = args[2]
+        padding = args[3] if len(args) >= 4 else (0, 0)
+        dilation = args[4] if len(args) >= 5 else (1, 1)
+        ceil_mode = args[5] if len(args) >= 6 else False
+
+        if not can_decompose_large_stride_maxpool2d(
+            kernel,
+            stride,
+            padding,
+            dilation,
+            ceil_mode,
+            x.data.shape,
+        ):
+            return super().call_operator(op, args, kwargs, meta)
+
+        kernel_h, kernel_w = (
+            (kernel, kernel) if isinstance(kernel, int) else kernel
+        )
+        n, c, height, width = x.data.shape
+        if isinstance(height, torch.SymInt) or isinstance(width, torch.SymInt):
+            return super().call_operator(op, args, kwargs, meta)
+
+        output_h = height // kernel_h
+        output_w = width // kernel_w
+        cropped_h = output_h * kernel_h
+        cropped_w = output_w * kernel_w
+
+        no_qparams_meta = meta.copy()
+        no_qparams_meta.data = meta.data.copy()
+        no_qparams_meta.data.pop("input_qparams", None)
+        no_qparams_meta.data.pop("output_qparams", None)
+
+        if cropped_h != height:
+            x = super().call_operator(
+                exir_ops.edge.aten.slice_copy.Tensor,
+                (x, 2, 0, cropped_h),
+                {},
+                no_qparams_meta,
+            )
+        if cropped_w != width:
+            x = super().call_operator(
+                exir_ops.edge.aten.slice_copy.Tensor,
+                (x, 3, 0, cropped_w),
+                {},
+                no_qparams_meta,
+            )
+
+        x = super().call_operator(
+            exir_ops.edge.aten.view_copy.default,
+            (x, [n, c, output_h, kernel_h, output_w, kernel_w]),
+            {},
+            no_qparams_meta,
+        )
+        x = super().call_operator(
+            exir_ops.edge.aten.permute_copy.default,
+            (x, [0, 1, 2, 4, 3, 5]),
+            {},
+            no_qparams_meta,
+        )
+        x = super().call_operator(
+            exir_ops.edge.aten.view_copy.default,
+            (x, [n, c, output_h * output_w * kernel_h, kernel_w]),
+            {},
+            no_qparams_meta,
+        )
+        x = super().call_operator(
+            op,
+            (x, (1, kernel_w), (1, 1), (0, 0), (1, 1), False),
+            {},
+            no_qparams_meta,
+        )
+        x = super().call_operator(
+            exir_ops.edge.aten.view_copy.default,
+            (x, [n, c, output_h * output_w, kernel_h]),
+            {},
+            no_qparams_meta,
+        )
+        x = super().call_operator(
+            op,
+            (x, (1, kernel_h), (1, 1), (0, 0), (1, 1), False),
+            {},
+            no_qparams_meta,
+        )
+        return super().call_operator(
+            exir_ops.edge.aten.view_copy.default,
+            (x, [n, c, output_h, output_w]),
+            {},
+            meta,
+        )

backends/arm/test/passes/test_decompose_large_stride_maxpool2d_pass.py

@@ -0,0 +1,81 @@
+# 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 types import SimpleNamespace
+from typing import Tuple
+from unittest.mock import patch
+
+import torch
+from executorch.backends.arm._passes.decompose_large_stride_maxpool2d_pass import (
+    DecomposeLargeStrideMaxPool2dPass,
+)
+from executorch.backends.arm._passes.remove_getitem_pass import RemoveGetItemPass
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+
+input_t = Tuple[torch.Tensor]
+
+
+class MaxPool1d(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.max_pool1d(x, kernel_size=5, stride=5)
+
+
+class MaxPool2d(torch.nn.Module):
+    def __init__(
+        self,
+        kernel_size: int | tuple[int, int],
+        stride: int | tuple[int, int],
+    ) -> None:
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.max_pool2d(
+            x,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+        )
+
+
+def _run_pass(
+    module: torch.nn.Module, inputs: input_t, expected_pool_count: int
+) -> None:
+    pipeline = PassPipeline[input_t](
+        module,
+        inputs,
+        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": expected_pool_count,
+        },
+        pass_list=[RemoveGetItemPass, DecomposeLargeStrideMaxPool2dPass],
+    )
+    with patch(
+        "executorch.backends.arm._passes.decompose_large_stride_maxpool2d_pass.get_context_spec",
+        return_value=SimpleNamespace(is_U55_subset=True),
+    ):
+        pipeline.run()
+
+
+def test_decompose_large_stride_max_pool1d() -> None:
+    _run_pass(MaxPool1d(), (torch.randn(1, 3, 17),), 2)
+
+
+def test_decompose_large_square_stride_max_pool2d() -> None:
+    _run_pass(MaxPool2d((5, 5), (5, 5)), (torch.randn(1, 3, 13, 17),), 2)
+
+
+def test_decompose_large_rectangular_stride_max_pool2d() -> None:
+    _run_pass(MaxPool2d((4, 7), (4, 7)), (torch.randn(1, 2, 11, 23),), 2)
+
+
+def test_keep_overlapping_large_stride_max_pool2d() -> None:
+    _run_pass(MaxPool2d((6, 6), (4, 4)), (torch.randn(1, 2, 15, 15),), 1)
+
+
+def test_keep_supported_scalar_pool_attributes() -> None:
+    _run_pass(MaxPool2d(2, 2), (torch.randn(1, 2, 15, 15),), 1)

backends/arm/tosa/partitioner.py

@@ -1,3 +1,3 @@
 # Copyright 2023-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -22,6 +22,9 @@
 
 import torch
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
+from executorch.backends.arm._passes.decompose_large_stride_maxpool2d_pass import (
+    can_decompose_large_stride_maxpool2d,
+)
 from executorch.backends.arm._passes.convert_expand_copy_to_repeat import (
     calculate_multiples,
 )
@@ -53,6 +56,36 @@
 logger = logging.getLogger(__name__)
 
 
+class DecomposableLargeStrideMaxPool2dSupported(OperatorSupportBase):
+    """Accept U55 max-pool nodes that backend preprocessing can legalize."""
+
+    def __init__(self, tosa_spec: TosaSpecification) -> None:
+        self.tosa_spec = tosa_spec
+
+    def is_node_supported(
+        self,
+        submodules: Mapping[str, torch.nn.Module],
+        node: torch.fx.Node,
+    ) -> bool:
+        """Return True when backend preprocessing can legalize the max pool."""
+        del submodules
+        if not self.tosa_spec.is_U55_subset or node.target not in {
+            exir_ops.edge.aten.max_pool2d.default,
+            exir_ops.edge.aten.max_pool2d_with_indices.default,
+        }:
+            return False
+
+        input_shape = get_first_fake_tensor(node.all_input_nodes[0]).shape
+        return can_decompose_large_stride_maxpool2d(
+            node.args[1],
+            node.args[2],
+            node.args[3] if len(node.args) >= 4 else (0, 0),
+            node.args[4] if len(node.args) >= 5 else (1, 1),
+            node.args[5] if len(node.args) >= 6 else False,
+            input_shape,
+        )
+
+
 class DecomposableResizeSupported(OperatorSupportBase):
     """Accept exact boundary bilinear downscales.
 
@@ -572,7 +605,10 @@
             containing_program,
             reporter,
             self.additional_checks,
-            additional_positive_checks=[self._decomposable_resize_support],
+            additional_positive_checks=[
+                self._decomposable_resize_support,
+                DecomposableLargeStrideMaxPool2dSupported(self.tosa_spec),
+            ],
         )
 
     def partition(self, exported_program: ExportedProgram) -> PartitionResult: