diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index bce7adc6e7d..62ecdaec3d2 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -629,3 +629,43 @@ def _upsample_nearest2d_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) +from executorch.backends.webgpu.test.ops.test_max_pool2d import ( + _det_input as _maxpool_det_input, + MaxPool2dModule, +) + + +@register_op_test("max_pool2d") +def _max_pool2d_suite() -> WebGPUTestSuite: + # SAM2 Hiera q_pool (native shape) + real Hiera channel count (ch768) + + # a padding case + a tiny eyeball case. + return WebGPUTestSuite( + module_factory=lambda kernel_size, stride, padding: MaxPool2dModule( + kernel_size, stride, padding + ), + cases=[ + Case( + name="q_pool", + construct={"kernel_size": 2, "stride": 2, "padding": 0}, + inputs=(InputSpec(shape=(1, 8, 12, 12), gen=_maxpool_det_input),), + ), + Case( + name="ch768", + construct={"kernel_size": 2, "stride": 2, "padding": 0}, + inputs=(InputSpec(shape=(1, 768, 14, 14), gen=_maxpool_det_input),), + ), + Case( + name="pad1", + construct={"kernel_size": 3, "stride": 2, "padding": 1}, + inputs=(InputSpec(shape=(1, 8, 7, 7), gen=_maxpool_det_input),), + ), + Case( + name="tiny", + construct={"kernel_size": 2, "stride": 2, "padding": 0}, + inputs=(InputSpec(shape=(1, 4, 5, 5), gen=_maxpool_det_input),), + ), + ], + golden_dtype="float32", + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_max_pool2d.py b/backends/webgpu/test/ops/test_max_pool2d.py new file mode 100644 index 00000000000..523ed75d7ee --- /dev/null +++ b/backends/webgpu/test/ops/test_max_pool2d.py @@ -0,0 +1,42 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.max_pool2d_with_indices.default` module for the WebGPU op-test +framework. + +`F.max_pool2d` decomposes to `aten.max_pool2d_with_indices.default`, a +multi-output op whose out is a ValueList `[values, indices]`; the handler +writes the VALUES only and never the int64 indices +(`runtime/ops/max_pool2d/MaxPool2d.cpp`). Max-pool is on the SAM2 Hiera +q_pool path. +""" + +import torch + + +class MaxPool2dModule(torch.nn.Module): + def __init__(self, kernel_size: int, stride: int, padding: int) -> None: + super().__init__() + self.kernel_size = kernel_size + self.stride = stride + self.padding = padding + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.max_pool2d( + x, + kernel_size=self.kernel_size, + stride=self.stride, + padding=self.padding, + ) + + +def _det_input(shape) -> torch.Tensor: + # ((i % 23) - 11) / 16: exact in fp32, spans negatives through positives. + n = 1 + for s in shape: + n *= s + idx = torch.arange(n, dtype=torch.int64) + return (((idx % 23) - 11).to(torch.float32) / 16.0).reshape(shape)