diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index eb5b2589193..71df953a65b 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -68,6 +68,9 @@ DecomposeIndexTensorToGatherPass, ) from .decompose_int_pow_pass import DecomposeIntPowPass # noqa +from .decompose_large_stride_maxpool2d_pass import ( # noqa + DecomposeLargeStrideMaxPool2dForU55Pass, +) from .decompose_layernorm_pass import DecomposeLayerNormPass # noqa from .decompose_leaky_relu_pass import DecomposeLeakyReLUPass # noqa from .decompose_linalg_vector_norm_pass import DecomposeLinalgVectorNormPass # noqa diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 947125fbd52..e265a3a4dfa 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -69,6 +69,7 @@ DecomposeIndexSelectToGatherPass, DecomposeIndexTensorToGatherPass, DecomposeIntPowPass, + DecomposeLargeStrideMaxPool2dForU55Pass, DecomposeLayerNormPass, DecomposeLeakyReLUPass, DecomposeLinalgVectorNormPass, @@ -612,6 +613,7 @@ def _tosa_pipeline( DecomposeCumsumPass(exported_program), DecomposeAsStridedCopyPass(), DecomposeMaxPool2dPass(), + DecomposeLargeStrideMaxPool2dForU55Pass(), 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..f7a38d52b65 --- /dev/null +++ b/backends/arm/_passes/decompose_large_stride_maxpool2d_pass.py @@ -0,0 +1,206 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from collections.abc import Sequence +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 +_U55_MAX_POOL_DIM = 65536 +_U55_MAX_POOL_KERNEL_PRODUCT = 65536 +_U55_MAX_POOL_KERNEL_WIDTH = 256 + + +def _pair(value, fallback: tuple[int, int] | None = None) -> tuple[int, int]: + if value is None: + if fallback is None: + raise ValueError("fallback is required when value is None") + return fallback + if isinstance(value, int): + return (value, value) + if isinstance(value, Sequence): + if len(value) == 0: + if fallback is None: + raise ValueError("fallback is required when value is empty") + return fallback + if len(value) < 2: + raise ValueError("expected sequence pair") + return (value[0], value[1]) + raise TypeError(f"Expected int or sequence pair, got {type(value)}") + + +# Keep these local to avoid importing operator support during pass construction; +# the constraints mirror pool_2d_support.dim_check/kernel_check for U55. +def _u55_dim_check(shape) -> bool: + return all( + not isinstance(dim, torch.SymInt) and 1 <= dim <= _U55_MAX_POOL_DIM + for dim in shape[1:] + ) + + +def _u55_kernel_check(kernel: tuple[int, int]) -> bool: + return ( + 1 <= kernel[0] * kernel[1] <= _U55_MAX_POOL_KERNEL_PRODUCT + and 1 <= kernel[1] <= _U55_MAX_POOL_KERNEL_WIDTH + ) + + +def can_decompose_large_stride_maxpool2d( + kernel, + stride, + padding, + dilation, + ceil_mode, + input_shape, +) -> bool: + kernel_h, kernel_w = _pair(kernel) + stride_h, stride_w = _pair(stride, (kernel_h, kernel_w)) + padding_h, padding_w = _pair(padding, (0, 0)) + dilation_h, dilation_w = _pair(dilation, (1, 1)) + height, width = input_shape[-2:] + + if ( + isinstance(height, torch.SymInt) + or isinstance(width, torch.SymInt) + or not _u55_kernel_check((kernel_h, kernel_w)) + or not _u55_kernel_check((1, kernel_w)) + or not _u55_kernel_check((1, kernel_h)) + or height < kernel_h + or width < kernel_w + ): + return False + + output_h = height // kernel_h + output_w = width // kernel_w + first_reduction_shape = ( + *input_shape[:-2], + output_h * output_w * kernel_h, + kernel_w, + ) + second_reduction_shape = (*input_shape[:-2], output_h * output_w, kernel_h) + output_shape = (*input_shape[:-2], output_h, output_w) + + return ( + 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 _u55_dim_check(input_shape) + and _u55_dim_check(first_reduction_shape) + and _u55_dim_check(second_reduction_shape) + and _u55_dim_check(output_shape) + ) + + +class DecomposeLargeStrideMaxPool2dForU55Pass(ArmOpTargetedPass): + """Legalize non-overlapping max_pool2d with strides unsupported by U55. + + Non-U55 profiles, including U85, use the normal TOSA/Vela path and do not + need this U55 pooling-engine workaround. + + """ + + _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] if len(args) >= 3 else kernel + 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 = _pair(kernel) + n, c, height, width = x.data.shape + 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..d51c9506372 --- /dev/null +++ b/backends/arm/test/passes/test_decompose_large_stride_maxpool2d_pass.py @@ -0,0 +1,174 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# 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 ( + can_decompose_large_stride_maxpool2d, + DecomposeLargeStrideMaxPool2dForU55Pass, +) +from executorch.backends.arm._passes.remove_getitem_pass import RemoveGetItemPass +from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, + PassPipeline, +) +from executorch.exir import EdgeCompileConfig, to_edge + +input_t = Tuple[torch.Tensor] + +_GET_CONTEXT_SPEC_PATCH = ( + "executorch.backends.arm._passes.decompose_large_stride_maxpool2d_pass." + "get_context_spec" +) + + +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, + ) -> 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 _transformed_module( + module: torch.nn.Module, inputs: input_t, is_u55: bool +) -> torch.nn.Module: + exported_program = torch.export.export(module.eval(), inputs, strict=True) + edge_program = to_edge( + exported_program, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + with patch( + _GET_CONTEXT_SPEC_PATCH, + return_value=SimpleNamespace(is_U55_subset=is_u55), + ): + transformed = edge_program.transform( + [RemoveGetItemPass(), DecomposeLargeStrideMaxPool2dForU55Pass()] + ) + return transformed.exported_program().module() + + +def _assert_pass_matches_eager( + module: torch.nn.Module, inputs: input_t, is_u55: bool = True +) -> None: + torch.testing.assert_close( + module.eval()(*inputs), + _transformed_module(module, inputs, is_u55)(*inputs), + ) + + +def _run_pass( + module: torch.nn.Module, + inputs: input_t, + expected_pool_count: int, + is_u55: bool = True, +) -> 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, DecomposeLargeStrideMaxPool2dForU55Pass], + ) + with patch( + _GET_CONTEXT_SPEC_PATCH, + return_value=SimpleNamespace(is_U55_subset=is_u55), + ): + pipeline.run() + _assert_pass_matches_eager(module, inputs, is_u55) + + +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_decompose_large_stride_max_pool2d_with_default_stride() -> None: + _run_pass(MaxPool2d((5, 5), None), (torch.randn(1, 3, 13, 17),), 2) + + +def test_keep_large_stride_max_pool2d_for_non_u55() -> None: + _run_pass( + MaxPool2d((5, 5), (5, 5)), + (torch.randn(1, 3, 13, 17),), + 1, + is_u55=False, + ) + + +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) + + +def test_reject_pool_exceeding_u55_original_dim_limit() -> None: + assert not can_decompose_large_stride_maxpool2d( + (5, 5), + (5, 5), + (0, 0), + (1, 1), + False, + (1, 1, 65537, 25), + ) + + +def test_reject_pool_exceeding_u55_intermediate_dim_limit() -> None: + assert not can_decompose_large_stride_maxpool2d( + (4, 4), + (4, 4), + (0, 0), + (1, 1), + False, + (1, 1, 1028, 1028), + ) + + +def test_decompose_large_stride_max_pool2d_u55_INT_pipeline() -> None: + inputs = (torch.randn(1, 3, 13, 17),) + pipeline = EthosU55PipelineINT[input_t]( + MaxPool2d((5, 5), (5, 5)), + inputs, + [], + [], + run_on_fvp=False, + ) + pipeline.pop_stage("check_not.exir") + pipeline.pop_stage("check_count.exir") + pipeline.pop_stage("to_executorch") + pipeline.run() + _assert_pass_matches_eager(MaxPool2d((5, 5), (5, 5)), inputs) diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index 09836a2121e..00f26cab6f3 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -25,6 +25,9 @@ from executorch.backends.arm._passes.convert_expand_copy_to_repeat import ( calculate_multiples, ) +from executorch.backends.arm._passes.decompose_large_stride_maxpool2d_pass import ( + can_decompose_large_stride_maxpool2d, +) from executorch.backends.arm._passes.decompose_unsupported_bilinear_resize_pass import ( is_exact_tosa_boundary_bilinear_downscale, ) @@ -53,6 +56,45 @@ logger = logging.getLogger(__name__) +class DecomposableLargeStrideMaxPool2dForU55Supported(OperatorSupportBase): + """Accept U55 max-pool nodes that backend preprocessing can legalize. + + Non-U55 profiles, including U85, use the standard max-pool support path. + This positive check is only for the U55 stride > 3 workaround. + + """ + + 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 + # The with_indices form is accepted because RemoveGetItemPass runs + # before backend preprocessing and canonicalizes value-only users to + # max_pool2d.default. If indices are used, RemoveGetItemPass rejects + # the node, matching the existing MaxPool2dSupported contract. + 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] if len(node.args) >= 3 else node.args[1], + 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 +614,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, + DecomposableLargeStrideMaxPool2dForU55Supported(self.tosa_spec), + ], ) def partition(self, exported_program: ExportedProgram) -> PartitionResult: