Skip to content

Commit e54d663

Browse files
apullinapullin
authored
Legalize large-stride non-overlapping maxpool on U55 (#21284)
Differential Revision: D113087299 Pull Request resolved: #21284
1 parent efd6b55 commit e54d663

5 files changed

Lines changed: 431 additions & 1 deletion

File tree

‎backends/arm/_passes/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,9 @@
6868
DecomposeIndexTensorToGatherPass,
6969
)
7070
from .decompose_int_pow_pass import DecomposeIntPowPass # noqa
71+
from .decompose_large_stride_maxpool2d_pass import ( # noqa
72+
DecomposeLargeStrideMaxPool2dForU55Pass,
73+
)
7174
from .decompose_layernorm_pass import DecomposeLayerNormPass # noqa
7275
from .decompose_leaky_relu_pass import DecomposeLeakyReLUPass # noqa
7376
from .decompose_linalg_vector_norm_pass import DecomposeLinalgVectorNormPass # noqa

‎backends/arm/_passes/arm_pass_manager.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@
6969
DecomposeIndexSelectToGatherPass,
7070
DecomposeIndexTensorToGatherPass,
7171
DecomposeIntPowPass,
72+
DecomposeLargeStrideMaxPool2dForU55Pass,
7273
DecomposeLayerNormPass,
7374
DecomposeLeakyReLUPass,
7475
DecomposeLinalgVectorNormPass,
@@ -612,6 +613,7 @@ def _tosa_pipeline(
612613
DecomposeCumsumPass(exported_program),
613614
DecomposeAsStridedCopyPass(),
614615
DecomposeMaxPool2dPass(),
616+
DecomposeLargeStrideMaxPool2dForU55Pass(),
615617
SizeAdjustInputPass(),
616618
DecomposeUnsupportedBilinearResizePass(self.tosa_spec),
617619
RewriteAdaptiveAvgPool2dPass(),

‎backends/arm/_passes/decompose_large_stride_maxpool2d_pass.py‎

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
from collections.abc import Sequence
8+
from typing import Set, Type
9+
10+
import torch
11+
from executorch.backends.arm._passes import ArmOpTargetedPass
12+
from executorch.backends.arm._passes.size_adjust_input_pass import SizeAdjustInputPass
13+
from executorch.backends.arm.tosa.specification import get_context_spec
14+
from executorch.exir.dialects._ops import ops as exir_ops
15+
from executorch.exir.pass_base import ExportPass
16+
17+
18+
_U55_MAX_POOL_STRIDE = 3
19+
_U55_MAX_POOL_DIM = 65536
20+
_U55_MAX_POOL_KERNEL_PRODUCT = 65536
21+
_U55_MAX_POOL_KERNEL_WIDTH = 256
22+
23+
24+
def _pair(value, fallback: tuple[int, int] | None = None) -> tuple[int, int]:
25+
if value is None:
26+
if fallback is None:
27+
raise ValueError("fallback is required when value is None")
28+
return fallback
29+
if isinstance(value, int):
30+
return (value, value)
31+
if isinstance(value, Sequence):
32+
if len(value) == 0:
33+
if fallback is None:
34+
raise ValueError("fallback is required when value is empty")
35+
return fallback
36+
if len(value) < 2:
37+
raise ValueError("expected sequence pair")
38+
return (value[0], value[1])
39+
raise TypeError(f"Expected int or sequence pair, got {type(value)}")
40+
41+
42+
# Keep these local to avoid importing operator support during pass construction;
43+
# the constraints mirror pool_2d_support.dim_check/kernel_check for U55.
44+
def _u55_dim_check(shape) -> bool:
45+
return all(
46+
not isinstance(dim, torch.SymInt) and 1 <= dim <= _U55_MAX_POOL_DIM
47+
for dim in shape[1:]
48+
)
49+
50+
51+
def _u55_kernel_check(kernel: tuple[int, int]) -> bool:
52+
return (
53+
1 <= kernel[0] * kernel[1] <= _U55_MAX_POOL_KERNEL_PRODUCT
54+
and 1 <= kernel[1] <= _U55_MAX_POOL_KERNEL_WIDTH
55+
)
56+
57+
58+
def can_decompose_large_stride_maxpool2d(
59+
kernel,
60+
stride,
61+
padding,
62+
dilation,
63+
ceil_mode,
64+
input_shape,
65+
) -> bool:
66+
kernel_h, kernel_w = _pair(kernel)
67+
stride_h, stride_w = _pair(stride, (kernel_h, kernel_w))
68+
padding_h, padding_w = _pair(padding, (0, 0))
69+
dilation_h, dilation_w = _pair(dilation, (1, 1))
70+
height, width = input_shape[-2:]
71+
72+
if (
73+
isinstance(height, torch.SymInt)
74+
or isinstance(width, torch.SymInt)
75+
or not _u55_kernel_check((kernel_h, kernel_w))
76+
or not _u55_kernel_check((1, kernel_w))
77+
or not _u55_kernel_check((1, kernel_h))
78+
or height < kernel_h
79+
or width < kernel_w
80+
):
81+
return False
82+
83+
output_h = height // kernel_h
84+
output_w = width // kernel_w
85+
first_reduction_shape = (
86+
*input_shape[:-2],
87+
output_h * output_w * kernel_h,
88+
kernel_w,
89+
)
90+
second_reduction_shape = (*input_shape[:-2], output_h * output_w, kernel_h)
91+
output_shape = (*input_shape[:-2], output_h, output_w)
92+
93+
return (
94+
max(stride_h, stride_w) > _U55_MAX_POOL_STRIDE
95+
and (kernel_h, kernel_w) == (stride_h, stride_w)
96+
and (padding_h, padding_w) == (0, 0)
97+
and (dilation_h, dilation_w) == (1, 1)
98+
and not ceil_mode
99+
and _u55_dim_check(input_shape)
100+
and _u55_dim_check(first_reduction_shape)
101+
and _u55_dim_check(second_reduction_shape)
102+
and _u55_dim_check(output_shape)
103+
)
104+
105+
106+
class DecomposeLargeStrideMaxPool2dForU55Pass(ArmOpTargetedPass):
107+
"""Legalize non-overlapping max_pool2d with strides unsupported by U55.
108+
109+
Non-U55 profiles, including U85, use the normal TOSA/Vela path and do not
110+
need this U55 pooling-engine workaround.
111+
112+
"""
113+
114+
_passes_required_after: Set[Type[ExportPass]] = {SizeAdjustInputPass}
115+
target_ops = (exir_ops.edge.aten.max_pool2d.default,)
116+
117+
def call_operator(self, op, args, kwargs, meta):
118+
if op not in self.target_ops or not get_context_spec().is_U55_subset:
119+
return super().call_operator(op, args, kwargs, meta)
120+
121+
x = args[0]
122+
kernel = args[1]
123+
stride = args[2] if len(args) >= 3 else kernel
124+
padding = args[3] if len(args) >= 4 else (0, 0)
125+
dilation = args[4] if len(args) >= 5 else (1, 1)
126+
ceil_mode = args[5] if len(args) >= 6 else False
127+
128+
if not can_decompose_large_stride_maxpool2d(
129+
kernel,
130+
stride,
131+
padding,
132+
dilation,
133+
ceil_mode,
134+
x.data.shape,
135+
):
136+
return super().call_operator(op, args, kwargs, meta)
137+
138+
kernel_h, kernel_w = _pair(kernel)
139+
n, c, height, width = x.data.shape
140+
output_h = height // kernel_h
141+
output_w = width // kernel_w
142+
cropped_h = output_h * kernel_h
143+
cropped_w = output_w * kernel_w
144+
145+
no_qparams_meta = meta.copy()
146+
no_qparams_meta.data = meta.data.copy()
147+
no_qparams_meta.data.pop("input_qparams", None)
148+
no_qparams_meta.data.pop("output_qparams", None)
149+
150+
if cropped_h != height:
151+
x = super().call_operator(
152+
exir_ops.edge.aten.slice_copy.Tensor,
153+
(x, 2, 0, cropped_h),
154+
{},
155+
no_qparams_meta,
156+
)
157+
if cropped_w != width:
158+
x = super().call_operator(
159+
exir_ops.edge.aten.slice_copy.Tensor,
160+
(x, 3, 0, cropped_w),
161+
{},
162+
no_qparams_meta,
163+
)
164+
165+
x = super().call_operator(
166+
exir_ops.edge.aten.view_copy.default,
167+
(x, [n, c, output_h, kernel_h, output_w, kernel_w]),
168+
{},
169+
no_qparams_meta,
170+
)
171+
x = super().call_operator(
172+
exir_ops.edge.aten.permute_copy.default,
173+
(x, [0, 1, 2, 4, 3, 5]),
174+
{},
175+
no_qparams_meta,
176+
)
177+
x = super().call_operator(
178+
exir_ops.edge.aten.view_copy.default,
179+
(x, [n, c, output_h * output_w * kernel_h, kernel_w]),
180+
{},
181+
no_qparams_meta,
182+
)
183+
x = super().call_operator(
184+
op,
185+
(x, (1, kernel_w), (1, 1), (0, 0), (1, 1), False),
186+
{},
187+
no_qparams_meta,
188+
)
189+
x = super().call_operator(
190+
exir_ops.edge.aten.view_copy.default,
191+
(x, [n, c, output_h * output_w, kernel_h]),
192+
{},
193+
no_qparams_meta,
194+
)
195+
x = super().call_operator(
196+
op,
197+
(x, (1, kernel_h), (1, 1), (0, 0), (1, 1), False),
198+
{},
199+
no_qparams_meta,
200+
)
201+
return super().call_operator(
202+
exir_ops.edge.aten.view_copy.default,
203+
(x, [n, c, output_h, output_w]),
204+
{},
205+
meta,
206+
)

0 commit comments

Comments
 (0)