From f1e2d983ca4ee7af5794cb68318cb61cfbaf8d71 Mon Sep 17 00:00:00 2001 From: Andrew Pullin Date: Fri, 24 Jul 2026 08:19:45 -0700 Subject: [PATCH] Legalize large-stride non-overlapping maxpool on U55 Summary: Teach the Arm partitioner to recognize exact U55 large-stride non-overlapping max-pool operations as safely decomposable, without mutating the graph during partitioning. Backend preprocessing then canonicalizes value-only max_pool2d_with_indices and rewrites eligible pools into supported stride-1 reductions. This preserves the partitioner contract, keeps unrelated TOSA/VGF/U85 paths unchanged, and also covers exported MaxPool1d. Differential Revision: D113087299 --- backends/arm/_passes/__init__.py | 3 + backends/arm/_passes/arm_pass_manager.py | 2 + .../decompose_large_stride_maxpool2d_pass.py | 152 ++++++++++++++++++ ...t_decompose_large_stride_maxpool2d_pass.py | 81 ++++++++++ backends/arm/tosa/partitioner.py | 38 ++++- 5 files changed, 275 insertions(+), 1 deletion(-) create mode 100644 backends/arm/_passes/decompose_large_stride_maxpool2d_pass.py create mode 100644 backends/arm/test/passes/test_decompose_large_stride_maxpool2d_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index eb5b2589193..11699ac2fbd 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -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 diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index da31e6e82d0..e122b3a7706 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -69,6 +69,7 @@ DecomposeIndexTensorToGatherPass, DecomposeIntPowPass, DecomposeLayerNormPass, + DecomposeLargeStrideMaxPool2dPass, DecomposeLeakyReLUPass, DecomposeLinalgVectorNormPass, DecomposeLinearPass, @@ -611,6 +612,7 @@ def _tosa_pipeline( DecomposeCumsumPass(exported_program), DecomposeAsStridedCopyPass(), DecomposeMaxPool2dPass(), + DecomposeLargeStrideMaxPool2dPass(), SizeAdjustInputPass(), DecomposeUnsupportedBilinearResizePass(self.tosa_spec), RewriteAdaptiveAvgPool2dPass(), diff --git a/backends/arm/_passes/decompose_large_stride_maxpool2d_pass.py b/backends/arm/_passes/decompose_large_stride_maxpool2d_pass.py new file mode 100644 index 00000000000..50b8ef24777 --- /dev/null +++ b/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, + ) diff --git a/backends/arm/test/passes/test_decompose_large_stride_maxpool2d_pass.py b/backends/arm/test/passes/test_decompose_large_stride_maxpool2d_pass.py new file mode 100644 index 00000000000..636cb78114a --- /dev/null +++ b/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) diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index 09836a2121e..270b25cc4e4 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -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 @@ def _create_operator_support( 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: