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

backends/xnnpack/operators/op_avg_pooling2d.py

@@ -4,7 +4,7 @@
 # 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 cast, Dict, List
+from typing import cast, Dict
 
 import torch
 from executorch.backends.xnnpack.operators.node_visitor import (
@@ -16,6 +16,7 @@
     XNNGraph,
     XNode,
 )
+from executorch.backends.xnnpack.utils.utils import normalize_pool2d_args
 from executorch.backends.xnnpack.utils.xnnpack_constants import XNN_FLAG_KEEP_DIMS
 
 
@@ -43,16 +44,10 @@ def define_node(
         # output
         output_id = vals_to_ids[node]
 
-        # kernel_size
-        pooling_height, pooling_width = cast(List, node.args[1])
-
-        # stride
-        stride_height, stride_width = cast(List, node.args[2])
-
-        # padding
-        padding_height, padding_width = 0, 0
-        if node.args[3] is not None:
-            padding_height, padding_width = cast(List[int], node.args[3])
+        kernel, stride, padding, _ = normalize_pool2d_args(node, has_dilation=False)
+        pooling_height, pooling_width = kernel
+        stride_height, stride_width = stride
+        padding_height, padding_width = padding
 
         ser_node = XNode(
             xnode_union=XNNAvgPooling2d(

backends/xnnpack/operators/op_max_pool2d.py

@@ -6,7 +6,7 @@
 
 # pyre-unsafe
 
-from typing import cast, Dict, List
+from typing import Dict
 
 import torch
 from executorch.backends.xnnpack.operators.node_visitor import (
@@ -20,6 +20,7 @@
     XNNMaxPooling2d,
     XNode,
 )
+from executorch.backends.xnnpack.utils.utils import normalize_pool2d_args
 from executorch.backends.xnnpack.utils.xnnpack_constants import XNN_FLAG_KEEP_DIMS
 
 
@@ -52,36 +53,23 @@ def define_node(
         kwargs["input_id"] = vals_to_ids[node.all_input_nodes[0]]
         kwargs["output_id"] = vals_to_ids[node]
 
-        # kernel info
-        kernel_shape = cast(List[int], node.args[1])
+        kernel_shape, stride, padding_shape, dilation = normalize_pool2d_args(
+            node, has_dilation=True
+        )
+
         kwargs["pooling_height"] = kernel_shape[0]
         kwargs["pooling_width"] = kernel_shape[1]
 
-        # stride info
-        stride = cast(List[int], node.args[2])
         kwargs["stride_height"] = stride[0]
         kwargs["stride_width"] = stride[1]
 
-        # padding info
-        kwargs["padding_top"] = 0
-        kwargs["padding_right"] = 0
-        kwargs["padding_bottom"] = 0
-        kwargs["padding_left"] = 0
-
-        if len(node.args) > 3:
-            padding_shape = cast(List[int], node.args[3])
-            kwargs["padding_top"] = padding_shape[0]
-            kwargs["padding_right"] = padding_shape[1]
-            kwargs["padding_bottom"] = padding_shape[0]
-            kwargs["padding_left"] = padding_shape[1]
-
-        # dilation info
-        kwargs["dilation_height"] = 1
-        kwargs["dilation_width"] = 1
-        if len(node._args) > 4:
-            dilation = cast(List[int], node.args[4])
-            kwargs["dilation_height"] = dilation[0]
-            kwargs["dilation_width"] = dilation[1]
+        kwargs["padding_top"] = padding_shape[0]
+        kwargs["padding_right"] = padding_shape[1]
+        kwargs["padding_bottom"] = padding_shape[0]
+        kwargs["padding_left"] = padding_shape[1]
+
+        kwargs["dilation_height"] = dilation[0]
+        kwargs["dilation_width"] = dilation[1]
 
         # ceil mode
         ceil_mode = node.args[5] if len(node.args) > 5 else False

backends/xnnpack/partition/config/generic_node_configs.py

@@ -21,7 +21,11 @@
     is_quant,
     tag_as_implicit_q_dq,
 )
-from executorch.backends.xnnpack.utils.utils import get_input_node, normalize_mean_dims
+from executorch.backends.xnnpack.utils.utils import (
+    get_input_node,
+    normalize_mean_dims,
+    normalize_pool2d_args,
+)
 from executorch.exir.backend.canonical_partitioners.config_partitioner import (
     format_target_name,
 )
@@ -180,7 +184,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 = cast(List[int], args[1])
+        kernel_size, _, _, _ = 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:
@@ -348,16 +352,15 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         if not self.check_common_constraints(node, ep):
             return False
 
-        kernel_size = node.args[1]
-        stride = node.args[2]
+        kernel_size, stride, _, _ = normalize_pool2d_args(node, has_dilation=True)
         is_ceil_mode = len(node.args) >= 6 and cast(bool, node.args[5])
 
         # Ceil mode is supported via op padding, which must be statically known.
         if is_ceil_mode and is_shape_dynamic(node):
             why(node, reason="ceil mode is not supported for dynamic shapes")
             return False
 
-        if stride[0] > kernel_size[0] or stride[1] > kernel_size[1]:  # pyre-ignore[16]
+        if stride[0] > kernel_size[0] or stride[1] > kernel_size[1]:
             why(
                 node,
                 reason=f"stride ({stride}) must be less than or equal to kernel size ({kernel_size})",

backends/xnnpack/test/ops/test_avgpool2d.py

@@ -78,6 +78,48 @@ def test_fp32_avgpool2d_count_include_pad_unsupported(self):
             .check_not(["torch.ops.higher_order.executorch_call_delegate"])
         )
 
+    class AvgPool2dSingleElementKernel(torch.nn.Module):
+        def __init__(self, divisor_override=None):
+            super().__init__()
+            self.avgPool = torch.nn.AvgPool2d(
+                kernel_size=[3],
+                count_include_pad=False,
+                divisor_override=divisor_override,
+            )
+
+        def forward(self, x):
+            return self.avgPool(x)
+
+    def test_fp32_avgpool2d_single_element_kernel_size(self):
+        """
+        int[2] arguments accept a 1-element list, which broadcasts to both dims.
+        """
+        inputs = (torch.randn(1, 1, 10, 10),)
+        (
+            Tester(self.AvgPool2dSingleElementKernel(), 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})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+    def test_fp32_avgpool2d_single_element_kernel_size_divisor_override(self):
+        """
+        A broadcast kernel_size must still produce a real pooling region to compare
+        against divisor_override. Here the region is 3 * 3, so 5 is unsupported.
+        """
+        inputs = (torch.randn(1, 1, 10, 10),)
+        (
+            Tester(self.AvgPool2dSingleElementKernel(divisor_override=5), inputs)
+            .export()
+            .check_count({"torch.ops.aten.avg_pool2d.default": 1})
+            .to_edge_transform_and_lower()
+            .check_not(["torch.ops.higher_order.executorch_call_delegate"])
+        )
+
     def test_fp32_avgpool2d_divisor_override(self):
         """
         The XNNPACK backend does not support divisor overrides not equal to the pooling region.

backends/xnnpack/test/ops/test_maxpool2d.py

@@ -188,6 +188,43 @@ def test_fp32_maxpool2d_unsupported_stride(self):
             .run_method_and_compare_outputs()
         )
 
+    def test_fp32_maxpool2d_default_stride(self):
+        """
+        stride defaults to kernel_size when omitted, and is absent from node.args.
+        """
+
+        class MaxPool2dDefaultStride(torch.nn.Module):
+            def forward(self, x):
+                return torch.nn.functional.max_pool2d(x, 2)
+
+        inputs = (torch.randn(4, 3, 24, 24),)
+        (
+            Tester(MaxPool2dDefaultStride(), inputs)
+            .export()
+            .check_count({"torch.ops.aten.max_pool2d.default": 1})
+            .to_edge_transform_and_lower()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+    def test_fp32_maxpool2d_single_element_kernel_size(self):
+        """
+        int[2] arguments accept a 1-element list, which broadcasts to both dims.
+        """
+        inputs = (torch.randn(4, 3, 24, 24),)
+        (
+            Tester(self.MaxPool2d(kernel_size=[3], stride=[2]), inputs)
+            .export()
+            .check_count({"torch.ops.aten.max_pool2d.default": 1})
+            .to_edge_transform_and_lower()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
     def test_qs8_maxpool2d(self):
         class MaxPool(torch.nn.Module):
             def __init__(self, maxpool_params):

backends/xnnpack/utils/utils.py

@@ -74,6 +74,33 @@ def normalize_mean_dims(mean_dims: Sequence[int] | int | None, rank: int) -> Lis
     return normalized_dims
 
 
+def normalize_pool2d_args(
+    node: torch.fx.Node, has_dilation: bool
+) -> Tuple[List[int], List[int], List[int], List[int]]:
+    """Return (kernel, stride, padding, dilation) for a pool2d node as 2-element lists.
+
+    Applies the aten schema defaults: an int or 1-element list broadcasts to both
+    dimensions, an omitted or empty stride means kernel_size, padding defaults to
+    0 and dilation to 1.
+    """
+    args = node.args
+
+    def pair(index: int, default: Optional[List[int]] = None) -> List[int]:
+        value = args[index] if len(args) > index else None
+        if isinstance(value, int):
+            return [value, value]
+        values = list(value) if value is not None else []
+        if not values:
+            return cast(List[int], default)
+        return [values[0], values[0]] if len(values) == 1 else [values[0], values[1]]
+
+    kernel = pair(1)
+    stride = pair(2, kernel)
+    padding = pair(3, [0, 0])
+    dilation = pair(4, [1, 1]) if has_dilation else [1, 1]
+    return kernel, stride, padding, dilation
+
+
 def get_relu_fused_node(node: torch.fx.Node) -> Optional[torch.fx.Node]:
     """
     Checks if the current node is only consumed by a relu node and can be fused,