Fix XNNPACK pooling crashing on default stride and single-element kernel_size

[slipstr34m]

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:

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

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:

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:

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

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21489

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 75962fe with merge base a1d4d35 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: slipstr34m / name: Shresth Rana (75962fe)

[nil-is-all]

Thanks for filing the issue and the PR, @slipstr34m. Running CI tests now.

[JakeStevens]

lgtm, waiting on CI results

[JakeStevens]

nit: these are all fixed sizes lists for the aten schema so we could consider immutable tuple instead

[slipstr34m]

On the tuple nit, agreed, and I have it ready locally. I held off pushing since a new commit re-triggers the CI approval, and that seemed like a poor trade for a nit on an already-green PR. Happy to push it if you'd rather have it in before merge.

[slipstr34m]

@JakeStevens CI finished green after your approval.