Fix XNNPACK pooling crashing on default stride and single-element kernel_size (#21489)
Fixes #21488
Fixes #10968
Fixes #10969
## Problem
The aten schemas are
```
avg_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, ...)
max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0,
int[2] dilation=1, bool ceil_mode=False)
```
Two schema facts the XNNPACK pooling code does not account for. `stride`
defaults to the empty list, meaning "use kernel_size", and may be absent
from `node.args` entirely. And `int[2]` accepts a scalar or a 1-element
list, which broadcasts to both dimensions.
`MaxPool2dConfig.check_constraints` indexed positionally:
```python
kernel_size = node.args[1]
stride = node.args[2]
...
if stride[0] > kernel_size[0] or stride[1] > kernel_size[1]:
```
and `AvgPoolingConfig.check_constraints` did
```python
kernel_size = cast(List[int], args[1])
pooling_region = kernel_size[0] * kernel_size[1]
```
Both assume a fully populated 2-element list, so three plain uses of the
documented API raise `IndexError` from inside the partitioner's own
constraint check, before any support decision is reached:
```
F.max_pool2d(x, 2) IndexError: tuple index out of range generic_node_configs.py:352
nn.MaxPool2d([3]) IndexError: list index out of range generic_node_configs.py:360
nn.AvgPool2d([3]) IndexError: list index out of range generic_node_configs.py:184
```
The visitors index the same way, so even a node that passed the config
would fail during serialization.
Declaring `target_name = "max_pool2d.default"` promises the whole
schema. The identical models with an explicit 2-element stride lower and
run correctly today.
## Fix
One helper, `normalize_pool2d_args`, placed next to the existing
`normalize_mean_dims` in `backends/xnnpack/utils/utils.py`, which exists
for exactly this reason on `mean.dim` and is already called from both
that config and that visitor:
```python
kernel = pair(1)
stride = pair(2, kernel)
padding = pair(3, [0, 0])
dilation = pair(4, [1, 1]) if has_dilation else [1, 1]
```
It is called from `AvgPoolingConfig.check_constraints`,
`MaxPool2dConfig.check_constraints`, `MaxPooling2d.define_node` and
`AveragePooling2d.define_node` in place of the raw indexing.
This is the same normalization five other backends here already do, so
"the partitioner may assume canonical args" is contradicted by standing
practice in the repo.
`backends/arm/_passes/rewrite_max_pool2d_pass.py:82-87` is structurally
the same helper:
```python
stride = to_2tuple(args[2]) if len(args) > 2 else ()
if not stride:
stride = kernel # default to kernel_size
padding = to_2tuple(args[3]) if len(args) > 3 else (0, 0)
dilation = to_2tuple(args[4]) if len(args) > 4 else (1, 1)
```
`backends/samsung/builders/op_max_pool2d.py:51-55` does the same,
including `stride = cast(List[int], node.args[2]) if len(node.args) > 2
else kernel_size`. `backends/mlx/ops.py` carries the comment `if not
stride: # empty list means default to kernel_size`.
`backends/qualcomm/builders/op_max_pool2d.py:63-76` and
`backends/apple/mps/operators/convolution_ops.py:53-57` both broadcast
1-element kernel, stride, padding and dilation.
`AvgPoolingConfig` already conceded the point internally. It guards
`len(args) >= 5`, `>= 6` and `>= 7` for `ceil_mode`, `count_include_pad`
and `divisor_override` on the lines around the crash. It knows arg
tuples arrive short, it just started guarding at index 5 instead of
index 2.
Two incidental cleanups caused by the change, called out so they are not
a surprise in review: the `# pyre-ignore[16]` on the stride comparison
is dropped because `stride` is now typed, and `cast`/`List` imports that
became unused are removed.
## Effect
Same script run on either side of the change. `delegated` is whether the
node reached XNNPACK, `maxdiff` is deviation from eager.
```
BEFORE AFTER
max_pool2d(x, 2) IndexError :352 delegated maxdiff=0.00e+00
nn.MaxPool2d([3]) #10969 IndexError :360 delegated maxdiff=0.00e+00
nn.AvgPool2d([3]) #10968 IndexError :184 delegated maxdiff=5.96e-08
CONTROL explicit [2,2] delegated maxdiff=0.00e+00 delegated maxdiff=0.00e+00
CONTROL explicit [3,3] delegated maxdiff=1.19e-07 delegated maxdiff=1.19e-07
NEGATIVE stride > kernel not delegated not delegated
```
The negative control is the important one. Normalization could have
turned "correctly decline" into "delegate something XNNPACK cannot do",
and it does not.
Declining these nodes instead of normalizing them would also have fixed
the crash, but it is strictly worse. All four cases match eager once
they delegate, so XNNPACK does support them, and silently dropping
`F.max_pool2d(x, 2)` out of the delegate would be a performance
regression on one of the most common CNN patterns there is.
## Blast radius
Only nodes whose pooling args are currently missing or scalar, which is
exactly the set that raises today. For fully specified 2-element args
the helper returns the same lists, so existing behaviour is unchanged.
No serialization format or runtime change; the `.pte` schema already
carries separate height and width fields.
## Testing
Four cases added, two per file. `test_maxpool2d.py` gets default stride
and a 1-element `kernel_size`. `test_avgpool2d.py` gets a 1-element
`kernel_size`, and a 1-element `kernel_size` with `divisor_override=5`.
The second matters because normalizing turns the `divisor_override !=
pooling_region` comparison from a crash into a live code path, and it
must still decline, since the region is 3 * 3 rather than 5.
Verified failing-first by reverting only the four source files and
keeping the tests: all four fail with `IndexError` at
`generic_node_configs.py:352`, `:360` and `:184`.
```
backends/xnnpack/test/ops/test_maxpool2d.py + test_avgpool2d.py
main : 12 passed
this PR : 16 passed
backends/xnnpack/test/ops/ (whole directory)
main : 222 passed, 4 subtests passed, 0 failed
this PR : 226 passed, 4 subtests passed, 0 failed
```
Which spellings are affected, for reference:
```
nn.MaxPool2d(2) kernel=[2, 2] stride=[2, 2] works today
nn.MaxPool2d(2, stride=2) kernel=[2, 2] stride=[2, 2] works today
F.max_pool2d(x, 2, 2) kernel=[2, 2] stride=[2, 2] works today
nn.MaxPool2d([3]) kernel=[3] stride=[3] IndexError
F.max_pool2d(x, 2) kernel=[2, 2] stride=MISSING IndexError
F.max_pool2d(x, [3]) kernel=[3] stride=MISSING IndexError
```
`nn.MaxPool2d` supplies `stride` from `kernel_size` in its own
`__init__`, and a bare int is expanded to `[n, n]` before the op is
called, so the common spelling is safe on both counts and existing
coverage never reaches these paths.
`lintrunner` reports no issues on all six changed files.
cc @GregoryComer @digantdesai @cbilgin @JakeStevens S
shresth rana committed
f287863976428d74e75dec42e0d97df38774aec9
Parent: 43cb2b2
Committed by GitHub <noreply@github.com>
on 7/31/2026, 1:40:55 AM