[ExecuTorch][WebGPU] max_pool2d op tests

backends/webgpu/test/op_tests/cases.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -629,3 +629,43 @@
         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,
+    )

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)