[ExecuTorch][WebGPU] upsample_nearest2d op tests

backends/webgpu/test/op_tests/cases.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -586,3 +586,46 @@
         atol=1e-4,
         rtol=1e-3,
     )
+from executorch.backends.webgpu.test.ops.test_upsample_nearest2d import (
+    _det_input as _upsample_det_input,
+    UpsampleNearest2dModule,
+)
+
+
+@register_op_test("upsample_nearest2d")
+def _upsample_nearest2d_suite() -> WebGPUTestSuite:
+    # SAM2 FPN 2x upsample chain (36->72->144) + a cheap tiny eyeball case.
+    return WebGPUTestSuite(
+        module_factory=lambda scale_factor: UpsampleNearest2dModule(scale_factor),
+        cases=[
+            Case(
+                name="fpn_36_72",
+                construct={"scale_factor": 2.0},
+                inputs=(InputSpec(shape=(1, 8, 36, 36), gen=_upsample_det_input),),
+            ),
+            Case(
+                name="fpn_72_144",
+                construct={"scale_factor": 2.0},
+                inputs=(InputSpec(shape=(1, 8, 72, 72), gen=_upsample_det_input),),
+            ),
+            Case(
+                name="tiny",
+                construct={"scale_factor": 2.0},
+                inputs=(InputSpec(shape=(1, 4, 5, 7), gen=_upsample_det_input),),
+            ),
+            Case(
+                # Non-integer, non-2x ratio (5->8): floor(oh*5/8) and the
+                # half-pixel-center formula round((oh+0.5)*5/8-0.5) diverge at
+                # oh=3,6 (1 vs 2, 3 vs 4) — locks in the legacy "nearest"
+                # formula (see upsample_nearest2d.wgsl) against a genuinely
+                # discriminating ratio, not just the 2x cases above where both
+                # formulas happen to agree.
+                name="non_2x_ratio",
+                construct={"scale_factor": 1.6},
+                inputs=(InputSpec(shape=(1, 2, 5, 5), gen=_upsample_det_input),),
+            ),
+        ],
+        golden_dtype="float32",
+        atol=1e-4,
+        rtol=1e-3,
+    )

backends/webgpu/test/ops/test_upsample_nearest2d.py

@@ -0,0 +1,33 @@
+# 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.upsample_nearest2d.vec` module for the WebGPU op-test framework.
+
+Upsample is on the SAM2/SAM3 pixel-decoder / FPN path
+(`F.interpolate(..., mode="nearest")`).
+"""
+
+import torch
+
+
+class UpsampleNearest2dModule(torch.nn.Module):
+    def __init__(self, scale_factor: float) -> None:
+        super().__init__()
+        self.scale_factor = scale_factor
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.interpolate(
+            x, scale_factor=self.scale_factor, mode="nearest"
+        )
+
+
+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)