diff --git a/backends/xnnpack/operators/op_avg_pooling2d.py b/backends/xnnpack/operators/op_avg_pooling2d.py index 94cd06cc08e..fb2c5b9404e 100644 --- a/backends/xnnpack/operators/op_avg_pooling2d.py +++ b/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( diff --git a/backends/xnnpack/operators/op_max_pool2d.py b/backends/xnnpack/operators/op_max_pool2d.py index 9ce734e6b53..a8c6d01e388 100644 --- a/backends/xnnpack/operators/op_max_pool2d.py +++ b/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 diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index 2f45a8bba04..c7a3be5f65f 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/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,8 +352,7 @@ 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. @@ -357,7 +360,7 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: 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})", diff --git a/backends/xnnpack/test/ops/test_avgpool2d.py b/backends/xnnpack/test/ops/test_avgpool2d.py index c1f149e5a93..23ca5f20e5a 100644 --- a/backends/xnnpack/test/ops/test_avgpool2d.py +++ b/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. diff --git a/backends/xnnpack/test/ops/test_maxpool2d.py b/backends/xnnpack/test/ops/test_maxpool2d.py index f82b27b09ec..15b2dcfc126 100644 --- a/backends/xnnpack/test/ops/test_maxpool2d.py +++ b/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): diff --git a/backends/xnnpack/utils/utils.py b/backends/xnnpack/utils/utils.py index bbdf71db5e5..94f75a042e0 100644 --- a/backends/xnnpack/utils/utils.py +++ b/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,