Skip to content

Commit 3da6123

Browse files
authored
Add MaxPool1D decomposition pass support (pytorch#21446)
Summary: Implement DecomposeMaxPool1dPass to enable MaxPool1D support on ARM backend by decomposing max_pool1d into unsqueeze_copy → max_pool2d → squeeze_copy. ## Implementation Strategy ### Decomposition Approach (Optimal for TOSA/Vela) The pass decomposes max_pool1d into max_pool2d via unsqueeze_copy/squeeze_copy operations: 1. unsqueeze_copy(dim=2): (N, C, L) → (N, C, 1, L) - add height dimension 2. max_pool2d: with adapted params [k]→[1,k], [s]→[1,s], [p]→[0,p], [d]→[1,d] 3. squeeze_copy(dims=[2]): (N, C, 1, L_out) → (N, C, L_out) - remove height dimension ### Why This Approach is Optimal 1. **unsqueeze_copy and squeeze_copy map to TOSA RESHAPE** which is zero-cost in Vela: - Classified as memory_only_ops (Reshape, Squeeze, ExpandDims, Identity) - Bypassed entirely when conditions met (NPU-produced, single consumer) - Tensor equivalence enables memory aliasing (same address) 2. **TFA Pipeline Placement (before quantization)**: - unsqueeze_copy.default is in _one_to_one_shared_input_qspec - squeeze_copy.dims is added to _one_to_one_shared_input_qspec - max_pool2d is in _one_to_one_shared_input_or_input_act_qspec - All get proper SharedQuantizationSpec from the annotator automatically 3. **Quantization Handling**: - Clear qparams on intermediate unsqueeze_copy and squeeze_copy ops (let annotator fill them) - Preserve original meta on max_pool2d for proper tracing - MAX_POOL2D doesn't need zero-point handling (unlike AVG_POOL2D) ### TOSA/Vela Constraints Validated - U55: Stride ≤3 ✓, Kernel ≤256x256 ✓ - U85: Extended stride support via accumulator save/restore - Dilation: Handled by separate DecomposeMaxPool2dPass if needed Reviewed By: 3l1 Differential Revision: D91760459 Pull Request resolved: pytorch#21446
1 parent 283297e commit 3da6123

6 files changed

Lines changed: 192 additions & 2 deletions

File tree

‎backends/arm/_passes/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@
9191
from .decompose_lstm_pass import DecomposeLstmPass # noqa
9292
from .decompose_masked_fill_pass import DecomposeMaskedFillPass # noqa
9393
from .decompose_matmul import DecomposeMatmulPass # noqa
94+
from .decompose_max_pool1d_pass import DecomposeMaxPool1dPass # noqa
9495
from .decompose_maxpool2d_with_dilation_pass import DecomposeMaxPool2dPass # noqa
9596
from .decompose_meandim_pass import DecomposeMeanDimPass # noqa
9697
from .decompose_ne_pass import DecomposeNotEqualPass # noqa

‎backends/arm/_passes/arm_pass_manager.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@
8282
DecomposeLstmPass,
8383
DecomposeMaskedFillPass,
8484
DecomposeMatmulPass,
85+
DecomposeMaxPool1dPass,
8586
DecomposeMaxPool2dPass,
8687
DecomposeMeanDimPass,
8788
DecomposeNotEqualPass,
@@ -681,6 +682,7 @@ def _tosa_pipeline(
681682
UnsqueezeBeforeRepeatPass(),
682683
DecomposeCumsumPass(exported_program),
683684
DecomposeAsStridedCopyPass(),
685+
DecomposeMaxPool1dPass(),
684686
NormalizeMaxPool2dInputRankPass(),
685687
DecomposeMaxPool2dPass(),
686688
DecomposeLargeStrideMaxPool2dForU55Pass(),
@@ -856,6 +858,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
856858
DecomposePowTensorTensorPass(self.tosa_spec, tfa_pass=True),
857859
DecomposeLinalgVectorNormPass(tfa_pass=True),
858860
DecomposeSqrtPass(tfa_pass=True),
861+
DecomposeMaxPool1dPass(tfa_pass=True),
859862
DecomposeSoftmaxPass(
860863
tfa_pass=True,
861864
),
Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
# Copyright 2026 Arm Limited and/or its affiliates.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
8+
from typing import List, Optional, Set, Type, Union
9+
10+
import torch
11+
from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass
12+
from executorch.backends.arm._passes.arm_pass_utils import meta_without_qparams
13+
from executorch.exir.pass_base import ExportPass
14+
15+
16+
def _normalize_to_list(
17+
value: Optional[Union[int, List[int], tuple]],
18+
default: Optional[List[int]] = None,
19+
) -> List[int]:
20+
"""Normalize parameter to list: handle None, int, tuple, list."""
21+
if value is None:
22+
if default is None:
23+
raise ValueError("Value cannot be None without a default")
24+
return default
25+
if isinstance(value, int):
26+
return [value]
27+
return list(value)
28+
29+
30+
class DecomposeMaxPool1dPass(ArmOpTargetedPass):
31+
"""Decomposes max_pool1d into max_pool2d via unsqueeze_copy/squeeze_copy
32+
operations.
33+
34+
This pass runs in transform_for_annotation (TFA) pipeline before quantization,
35+
ensuring proper quantization annotation for the decomposed ops.
36+
37+
Transformation:
38+
max_pool1d(x, kernel, stride, padding, dilation, ceil_mode)
39+
→ unsqueeze_copy(x, dim=2) # (N,C,L) → (N,C,1,L)
40+
→ max_pool2d(..., [1,k], [1,s], [0,p], [1,d], ceil_mode)
41+
→ squeeze_copy(..., dims=[2]) # (N,C,1,L') → (N,C,L')
42+
43+
"""
44+
45+
target_ops = (torch.ops.aten.max_pool1d.default,)
46+
47+
_passes_required_after: Set[Type[ExportPass]] = set()
48+
49+
def call_operator(self, op, args, kwargs, meta):
50+
if op != torch.ops.aten.max_pool1d.default or not self.allowed_to_transform(
51+
meta
52+
):
53+
return super().call_operator(op, args, kwargs, meta)
54+
55+
# Extract and normalize arguments
56+
x = args[0]
57+
kernel_size = _normalize_to_list(args[1])
58+
stride = _normalize_to_list(
59+
args[2] if len(args) > 2 else None,
60+
default=kernel_size, # stride defaults to kernel_size
61+
)
62+
padding = _normalize_to_list(args[3] if len(args) > 3 else 0)
63+
dilation = _normalize_to_list(args[4] if len(args) > 4 else 1)
64+
ceil_mode = args[5] if len(args) > 5 else False
65+
66+
# Step 1: Unsqueeze input from 3D to 4D at dim=2
67+
# (N, C, L) → (N, C, 1, L)
68+
unsqueeze_meta = meta_without_qparams(meta)
69+
x_4d = super().call_operator(
70+
torch.ops.aten.unsqueeze_copy.default,
71+
(x, 2),
72+
{},
73+
unsqueeze_meta,
74+
updated=True,
75+
)
76+
77+
# Step 2: Call max_pool2d with 2D parameters
78+
# kernel: [k] → [1, k], stride: [s] → [1, s]
79+
# padding: [p] → [0, p], dilation: [d] → [1, d]
80+
pooled = super().call_operator(
81+
torch.ops.aten.max_pool2d.default,
82+
(
83+
x_4d,
84+
[1] + kernel_size,
85+
[1] + stride,
86+
[0] + padding,
87+
[1] + dilation,
88+
ceil_mode,
89+
),
90+
{},
91+
meta,
92+
updated=True,
93+
)
94+
95+
# Step 3: Squeeze output back to 3D at dims=[2]
96+
# (N, C, 1, L') → (N, C, L')
97+
squeeze_meta = meta_without_qparams(meta)
98+
output = super().call_operator(
99+
torch.ops.aten.squeeze_copy.dims,
100+
(pooled, [2]),
101+
{},
102+
squeeze_meta,
103+
updated=True,
104+
)
105+
106+
return output

‎backends/arm/quantizer/quantization_annotator.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -575,6 +575,9 @@ def _get_fixed_qparams_qspec(
575575
torch.ops.aten.squeeze_copy.default,
576576
torch.ops.aten.squeeze_copy.dim,
577577
torch.ops.aten.squeeze_.dim,
578+
# DecomposeMaxPool1dPass emits squeeze_copy.dims as a view-like intermediate;
579+
# include here so it receives SharedQuantizationSpec from its input.
580+
torch.ops.aten.squeeze_copy.dims,
578581
torch.ops.aten.squeeze.dim,
579582
torch.ops.aten.squeeze.dims,
580583
torch.ops.aten.unbind.int,

‎backends/arm/test/ops/test_max_pool1d.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -148,9 +148,7 @@ def test_max_pool2d_vgf_no_quant(test_data: Callable):
148148

149149

150150
@common.parametrize("test_data", test_data_suite_all)
151-
@pytest.mark.xfail(reason="MaxPool1D not yet supported", strict=False)
152151
@common.SkipIfNoModelConverter
153-
@pytest.mark.xfail(reason="Quantized MaxPool1D not yet supported in VGF", strict=False)
154152
def test_max_pool2d_vgf_quant_decomposed(test_data: Callable):
155153
"""Test max_pool1d with VGF pipeline (quantized)."""
156154
test_data, model_params = test_data()
Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
# Copyright 2026 Arm Limited and/or its affiliates.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
8+
import torch
9+
from executorch.backends.arm._passes import ArmPassManager, DecomposeMaxPool1dPass
10+
from executorch.backends.arm.test import common
11+
12+
13+
class _MaxPool1d(torch.nn.Module):
14+
def __init__(self, kernel_size: int, stride: int, padding: int = 0) -> None:
15+
super().__init__()
16+
self.pool = torch.nn.MaxPool1d(kernel_size, stride, padding)
17+
18+
def forward(self, x: torch.Tensor) -> torch.Tensor:
19+
return self.pool(x)
20+
21+
22+
def test_decompose_max_pool1d_to_max_pool2d() -> None:
23+
for kernel_size, stride, padding in [(4, 2, 0), (3, 2, 1), (3, 1, 0)]:
24+
module = _MaxPool1d(kernel_size, stride, padding).eval()
25+
inputs = (torch.randn(1, 8, 32),)
26+
exported = torch.export.export(module, inputs, strict=True)
27+
28+
result = DecomposeMaxPool1dPass()(exported.graph_module)
29+
assert result is not None
30+
graph_module = result.graph_module
31+
graph_module.graph.lint()
32+
33+
targets = {
34+
node.target
35+
for node in graph_module.graph.nodes
36+
if node.op == "call_function"
37+
}
38+
assert torch.ops.aten.max_pool1d.default not in targets
39+
assert torch.ops.aten.max_pool2d.default in targets
40+
41+
actual = graph_module(*inputs)
42+
if isinstance(actual, tuple):
43+
(actual,) = actual
44+
torch.testing.assert_close(module(*inputs), actual)
45+
46+
47+
def test_pass_skips_graph_with_submodules_but_no_targets() -> None:
48+
"""Graphs with call_module nodes but no max_pool1d must not retrace.
49+
50+
ExportPass retracing raises on call_module nodes; the targeted pre-scan must
51+
skip such graphs (e.g. stateful LSTM prepare graphs) entirely.
52+
53+
"""
54+
graph = torch.fx.Graph()
55+
x = graph.placeholder("x")
56+
out = graph.call_module("lin", args=(x,))
57+
graph.output(out)
58+
graph_module = torch.fx.GraphModule({"lin": torch.nn.Linear(8, 8)}, graph)
59+
60+
result = DecomposeMaxPool1dPass()(graph_module)
61+
62+
assert result is not None
63+
assert not result.modified
64+
65+
66+
def test_max_pool1d_decomposed_in_tfa_pipeline() -> None:
67+
module = _MaxPool1d(4, 2).eval()
68+
inputs = (torch.randn(1, 16, 50),)
69+
exported = torch.export.export(module, inputs, strict=True)
70+
71+
transformed = ArmPassManager(
72+
common.get_tosa_compile_spec("TOSA-1.0+INT")
73+
).transform_for_annotation_pipeline(exported.graph_module)
74+
75+
targets = {
76+
node.target for node in transformed.graph.nodes if node.op == "call_function"
77+
}
78+
assert torch.ops.aten.max_pool1d.default not in targets
79+
assert torch.ops.aten.max_pool2d.default in targets

0 commit comments

Comments
 (0)