Fix XNNPACK rejecting avg_pool2d with default arguments (pytorch#21862)

## Problem

`AvgPoolingConfig.check_constraints` rejects every node with
`count_include_pad=True`:

```python
if count_include_pad:
    why(
        node,
        reason="zero-padding in the averaging calculation is not supported",
    )
    return False
```

But `count_include_pad=True` is the aten default, so a plain
`nn.AvgPool2d(2)` — the most common way to write the layer — never
delegates to XNNPACK. The delegate splits around it and the node falls
back to the portable op, and the `WhyNoPartition` log tells the user
their model has zero-padding it does not have:

```
WhyNoPartition: Node aten_avg_pool2d_default was not partitioned because zero-padding in the averaging calculation is not supported.
```

The flag only has semantics when padded elements exist: it controls
whether zeros from padding count toward the averaging divisor. With the
default `padding=0` there is nothing to include or exclude, so the
rejection excludes nodes XNNPACK handles correctly. Verified
numerically: with `padding=0`, eager outputs for
`count_include_pad=True` and `False` are bit-identical, and the
delegated result matches eager within ordinary float tolerance; with
`padding=1` they genuinely differ, which is the case the backend does
not implement.

## Fix

Reject `count_include_pad=True` only when the node has non-zero padding,
reading padding from the existing `normalize_pool2d_args` call:

```python
kernel_size, _, padding, _ = normalize_pool2d_args(node, has_dilation=False)
...
if count_include_pad and any(p != 0 for p in padding):
```

A no-op for every node the config currently accepts (those all have
`count_include_pad=False`), and every currently-rejected padded node is
still rejected, so the only behavior change is crash-to-delegate for the
default-argument case. The visitor needs no change: it serializes the
padding values themselves and never reads `count_include_pad`, and
XNNPACK's divisor convention is irrelevant when no padded elements
exist.

## Effect

Same model, `nn.AvgPool2d(2)` with all-default arguments, through
`to_edge_transform_and_lower` with the XNNPACK partitioner:

```
before: delegates=0  remaining aten.avg_pool2d.default in graph, WhyNoPartition blames zero-padding
after:  delegates=1  no avg_pool2d left in graph, output matches eager
guard:  padding=(1,1) with count_include_pad=True still not delegated
```

## Testing

New `test_fp32_avgpool2d_default_args` lowers `nn.AvgPool2d(2)` end to
end through serialize and compares outputs against eager. Verified
failing-first: with the tests kept and only the source reverted, exactly
that test fails and the other seven pass, including
`test_fp32_avgpool2d_count_include_pad_unsupported` (padding=(1,1)),
which still rejects under the fix.

backends/xnnpack/test/ops/test_avgpool2d.py: 7 -> 8 passed. Full
backends/xnnpack/test/ops/ suite: 229 passed on main, 230 passed with
this change, no regressions. The suite runs exclude test_prelu.py, whose
test_fp32_prelu_constant_weight_empty_decompositions_file_load segfaults
identically on unmodified main in my environment (prebuilt runtime,
macOS arm64), so it cannot observe this change either way. lintrunner
reports no issues on both changed files.


cc @GregoryComer @digantdesai @cbilgin @JakeStevens

Author: slipstr34m

backends/xnnpack/partition/config/generic_node_configs.py

@@ -168,7 +168,8 @@ class AvgPoolingConfig(GenericNodePartitionerConfig):
 
     def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         """
-        XNNPACK does not support ceil_mode = True and count_include_pad = True
+        XNNPACK does not support ceil_mode = True, or count_include_pad = True
+        when the node has non-zero padding.
         Additionally, we only support divisor_override if divisor_override = pooling region
         """
         if not self.check_common_constraints(node, ep):
@@ -184,7 +185,7 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         if len(args) >= 6:
             count_include_pad = cast(bool, args[5])
 
-        kernel_size, _, _, _ = normalize_pool2d_args(node, has_dilation=False)
+        kernel_size, _, padding, _ = normalize_pool2d_args(node, has_dilation=False)
         pooling_region = kernel_size[0] * kernel_size[1]
         divisor_override = pooling_region  # Default divisor is pooling_region
         if len(args) >= 7:
@@ -194,7 +195,7 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
             why(node, reason="ceil mode is not supported")
             return False
 
-        if count_include_pad:
+        if count_include_pad and any(p != 0 for p in padding):
             why(
                 node,
                 reason="zero-padding in the averaging calculation is not supported",

backends/xnnpack/test/ops/test_avgpool2d.py

@@ -67,7 +67,8 @@ def test_fp32_avgpool2d_ceil_mode_unsupported(self):
 
     def test_fp32_avgpool2d_count_include_pad_unsupported(self):
         """
-        The XNNPACK backend does not support count_include_pad=True.
+        The XNNPACK backend does not support count_include_pad=True with
+        non-zero padding.
         """
         inputs = (torch.randn(1, 1, 10, 10),)
         (
@@ -78,6 +79,32 @@ def test_fp32_avgpool2d_count_include_pad_unsupported(self):
             .check_not(["torch.ops.higher_order.executorch_call_delegate"])
         )
 
+    class AvgPool2dAllDefaults(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.avgPool = torch.nn.AvgPool2d(2)
+
+        def forward(self, x):
+            return self.avgPool(x)
+
+    def test_fp32_avgpool2d_default_args(self):
+        """
+        count_include_pad defaults to True, but with the default zero padding
+        there are no padded elements to count, so the node must still delegate.
+        """
+        inputs = (torch.randn(1, 1, 10, 10),)
+        (
+            Tester(self.AvgPool2dAllDefaults(), inputs)
+            .export()
+            .check_count({"torch.ops.aten.avg_pool2d.default": 1})
+            .to_edge_transform_and_lower()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .check_not(["executorch_exir_dialects_edge__ops_aten_avg_pool2d_default"])
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
     class AvgPool2dSingleElementKernel(torch.nn.Module):
         def __init__(self, divisor_override=None):
             super().__init__()