diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index c7a3be5f65f..8a769a8383e 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/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", diff --git a/backends/xnnpack/test/ops/test_avgpool2d.py b/backends/xnnpack/test/ops/test_avgpool2d.py index 23ca5f20e5a..909569daf8e 100644 --- a/backends/xnnpack/test/ops/test_avgpool2d.py +++ b/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__()