[ExecuTorch][WebGPU] conv2d op tests

backends/webgpu/test/op_tests/cases.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -387,3 +387,63 @@
         atol=1e-4,
         rtol=1e-3,
     )
+from executorch.backends.webgpu.test.ops.test_conv2d import (
+    _chw_ramp,
+    make_conv,
+)
+
+
+@register_op_test("conv2d")
+def _conv2d_suite() -> WebGPUTestSuite:
+    # DaViT patch-embed / downsample convs + conv_transpose2d (same registration,
+    # folded by the `transposed` arg). NCHW fp32.
+    return WebGPUTestSuite(
+        module_factory=make_conv,
+        cases=[
+            Case(
+                name="conv3x3_pad1",
+                construct={"in_ch": 8, "out_ch": 16, "kernel": 3, "padding": 1},
+                inputs=(InputSpec(shape=(1, 8, 16, 16), gen=_chw_ramp),),
+            ),
+            Case(
+                name="patch_embed",
+                construct={"in_ch": 3, "out_ch": 64, "kernel": 16, "stride": 16},
+                inputs=(InputSpec(shape=(1, 3, 32, 32), gen=_chw_ramp),),
+            ),
+            Case(
+                name="strided",
+                construct={
+                    "in_ch": 3,
+                    "out_ch": 8,
+                    "kernel": 3,
+                    "stride": 2,
+                    "padding": 1,
+                },
+                inputs=(InputSpec(shape=(1, 3, 16, 16), gen=_chw_ramp),),
+            ),
+            Case(
+                name="depthwise",
+                construct={
+                    "in_ch": 8,
+                    "out_ch": 8,
+                    "kernel": 3,
+                    "padding": 1,
+                    "groups": 8,
+                },
+                inputs=(InputSpec(shape=(1, 8, 8, 8), gen=_chw_ramp),),
+            ),
+            Case(
+                name="transpose2x",
+                construct={
+                    "in_ch": 4,
+                    "out_ch": 4,
+                    "kernel": 2,
+                    "stride": 2,
+                    "transposed": True,
+                },
+                inputs=(InputSpec(shape=(1, 4, 4, 4), gen=_chw_ramp),),
+            ),
+        ],
+        atol=1e-4,
+        rtol=1e-3,
+    )

backends/webgpu/test/ops/test_conv2d.py

@@ -0,0 +1,95 @@
+# 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.convolution.default` (conv2d + conv_transpose2d) module + inputs.
+
+`make_conv` and `_chw_ramp` are imported by `cases.py` to drive the declarative
+op-test suite. `Conv2dTest` is the export-delegation smoke test. conv2d is the
+DaViT patch-embed / downsample op in Florence-2; conv_transpose2d shares the
+same `aten.convolution.default` registration (folded by the `transposed` arg).
+"""
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+
+class Conv2dModule(torch.nn.Module):
+    def __init__(
+        self,
+        in_ch: int,
+        out_ch: int,
+        kernel: int,
+        stride: int = 1,
+        padding: int = 0,
+        groups: int = 1,
+        transposed: bool = False,
+    ):
+        super().__init__()
+        cls = torch.nn.ConvTranspose2d if transposed else torch.nn.Conv2d
+        self.conv = cls(
+            in_ch, out_ch, kernel, stride=stride, padding=padding, groups=groups
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+def make_conv(
+    in_ch: int,
+    out_ch: int,
+    kernel: int,
+    stride: int = 1,
+    padding: int = 0,
+    groups: int = 1,
+    transposed: bool = False,
+) -> torch.nn.Module:
+    """Factory with deterministic small weights/bias for reproducible goldens."""
+    m = Conv2dModule(in_ch, out_ch, kernel, stride, padding, groups, transposed)
+    with torch.no_grad():
+        w = m.conv.weight
+        w.copy_(
+            torch.linspace(-1.0, 1.0, w.numel(), dtype=torch.float32).reshape(w.shape)
+            / (kernel * kernel)
+        )
+        if m.conv.bias is not None:
+            b = m.conv.bias
+            b.copy_(torch.linspace(-0.5, 0.5, b.numel(), dtype=torch.float32))
+    return m
+
+
+def _chw_ramp(shape) -> torch.Tensor:
+    """Deterministic [N, C, H, W] ramp in [-1, 1]."""
+    n = 1
+    for d in shape:
+        n *= d
+    return torch.linspace(-1.0, 1.0, n, dtype=torch.float32).reshape(shape)
+
+
+def _export(m: torch.nn.Module, x: torch.Tensor):
+    ep = torch.export.export(m, (x,))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+class Conv2dTest(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        configs = [
+            ("conv", make_conv(3, 8, 3, padding=1), (1, 3, 8, 8)),
+            ("transpose", make_conv(4, 4, 2, stride=2, transposed=True), (1, 4, 4, 4)),
+        ]
+        for name, model, shape in configs:
+            et = _export(model.eval(), _chw_ramp(shape))
+            found = any(
+                d.id == "VulkanBackend"
+                for plan in et.executorch_program.execution_plan
+                for d in plan.delegates
+            )
+            self.assertTrue(found, f"Expected a VulkanBackend delegate (conv {name})")