Add MaxPool1D decomposition pass support (#17022) #17022
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,106 @@ | ||
| # Copyright (c) Meta Platforms, Inc. and affiliates. | ||
| # All rights reserved. | ||
| # Copyright 2025 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 List, Optional, Set, Type, Union | ||
|
|
||
| import torch | ||
| from executorch.backends.arm._passes.arm_pass import ArmPass | ||
| from executorch.exir.pass_base import ExportPass | ||
|
|
||
|
|
||
| def _normalize_to_list( | ||
| value: Optional[Union[int, List[int], tuple]], | ||
| default: Optional[List[int]] = None, | ||
| ) -> List[int]: | ||
| """Normalize parameter to list: handle None, int, tuple, list.""" | ||
| if value is None: | ||
| if default is None: | ||
| raise ValueError("Value cannot be None without a default") | ||
| return default | ||
| if isinstance(value, int): | ||
| return [value] | ||
| return list(value) | ||
|
|
||
|
|
||
| class DecomposeMaxPool1dPass(ArmPass): | ||
| """Decomposes max_pool1d into max_pool2d via unsqueeze_copy/squeeze_copy | ||
| operations. | ||
|
|
||
| This pass runs in transform_for_annotation (TFA) pipeline before quantization, | ||
| ensuring proper quantization annotation for the decomposed ops. | ||
|
|
||
| Transformation: | ||
| max_pool1d(x, kernel, stride, padding, dilation, ceil_mode) | ||
| → unsqueeze_copy(x, dim=2) # (N,C,L) → (N,C,1,L) | ||
| → max_pool2d(..., [1,k], [1,s], [0,p], [1,d], ceil_mode) | ||
| → squeeze_copy(..., dims=[2]) # (N,C,1,L') → (N,C,L') | ||
| """ | ||
|
Comment on lines
+29
to
+41
There was a problem hiding this comment. There are existing |
||
|
|
||
| _passes_required_after: Set[Type[ExportPass]] = set() | ||
|
|
||
| def call_operator(self, op, args, kwargs, meta): | ||
| if op != torch.ops.aten.max_pool1d.default or not self.allowed_to_transform( | ||
| meta | ||
| ): | ||
| return super().call_operator(op, args, kwargs, meta) | ||
|
Comment on lines
+45
to
+49
There was a problem hiding this comment.
|
||
|
|
||
| # Extract and normalize arguments | ||
| x = args[0] | ||
| kernel_size = _normalize_to_list(args[1]) | ||
| stride = _normalize_to_list( | ||
| args[2] if len(args) > 2 else None, | ||
| default=kernel_size, # stride defaults to kernel_size | ||
| ) | ||
| padding = _normalize_to_list(args[3] if len(args) > 3 else 0) | ||
| dilation = _normalize_to_list(args[4] if len(args) > 4 else 1) | ||
| ceil_mode = args[5] if len(args) > 5 else False | ||
|
Comment on lines
+54
to
+60
There was a problem hiding this comment. The default value handling for stride on line 56 may not be correct. When stride is not provided (None), it should default to kernel_size. However, if len(args) > 2 and args[2] is explicitly None, the current code will pass None to _normalize_to_list, which will use the default. This is correct. But if args[2] contains an empty list or tuple, this won't be handled correctly. Consider adding validation that kernel_size, stride, padding, and dilation after normalization all have exactly one element (for 1D operations), to catch any potential issues early. |
||
|
|
||
| # Step 1: Unsqueeze input from 3D to 4D at dim=2 | ||
| # (N, C, L) → (N, C, 1, L) | ||
| unsqueeze_meta = meta.copy() | ||
| unsqueeze_meta.data["input_qparams"] = {} | ||
| unsqueeze_meta.data["output_qparams"] = {} | ||
| x_4d = super().call_operator( | ||
| torch.ops.aten.unsqueeze_copy.default, | ||
| (x, 2), | ||
| {}, | ||
| unsqueeze_meta, | ||
| updated=True, | ||
| ) | ||
|
Comment on lines
+67
to
+73
There was a problem hiding this comment. When inserting the |
||
|
|
||
| # Step 2: Call max_pool2d with 2D parameters | ||
| # kernel: [k] → [1, k], stride: [s] → [1, s] | ||
| # padding: [p] → [0, p], dilation: [d] → [1, d] | ||
| pooled = super().call_operator( | ||
| torch.ops.aten.max_pool2d.default, | ||
| ( | ||
| x_4d, | ||
| [1] + kernel_size, | ||
| [1] + stride, | ||
| [0] + padding, | ||
| [1] + dilation, | ||
| ceil_mode, | ||
| ), | ||
| {}, | ||
| meta, | ||
| updated=True, | ||
| ) | ||
|
Comment on lines
+45
to
+91
|
||
|
|
||
| # Step 3: Squeeze output back to 3D at dims=[2] | ||
| # (N, C, 1, L') → (N, C, L') | ||
| squeeze_meta = meta.copy() | ||
| squeeze_meta.data["input_qparams"] = {} | ||
| squeeze_meta.data["output_qparams"] = {} | ||
| output = super().call_operator( | ||
| torch.ops.aten.squeeze_copy.dims, | ||
| (pooled, [2]), | ||
| {}, | ||
| squeeze_meta, | ||
| updated=True, | ||
| ) | ||
|
Comment on lines
+98
to
+104
There was a problem hiding this comment. Same as the inserted |
||
|
|
||
| return output | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Agree with this one. We shouldn't see any max_pool1d:s at this stage, the pass won't match any exir-targets and will emit torch.ops rather than exir_ops.