[ExecuTorch][WebGPU] addmm op tests

backends/webgpu/test/op_tests/cases.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -527,3 +527,22 @@
         atol=1e-4,
         rtol=1e-3,
     )
+from executorch.backends.webgpu.test.ops.test_addmm import (
+    _randn as _addmm_randn,
+    AddmmModule,
+)
+
+
+@register_op_test("addmm")
+def _addmm_suite() -> WebGPUTestSuite:
+    return WebGPUTestSuite(
+        module_factory=lambda n: AddmmModule(n),
+        cases=[
+            Case(name="small", construct={"n": 32}, inputs=(InputSpec(shape=(4, 16), gen=_addmm_randn), InputSpec(shape=(16, 32), gen=_addmm_randn))),
+            Case(name="bart", construct={"n": 768}, inputs=(InputSpec(shape=(16, 768), gen=_addmm_randn), InputSpec(shape=(768, 768), gen=_addmm_randn))),
+            Case(name="odd_k", construct={"n": 32}, inputs=(InputSpec(shape=(4, 15), gen=_addmm_randn), InputSpec(shape=(15, 32), gen=_addmm_randn))),
+        ],
+        golden_dtype="float32",
+        atol=1e-4,
+        rtol=1e-3,
+    )

backends/webgpu/test/ops/test_addmm.py

@@ -0,0 +1,28 @@
+# 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.addmm.default` module for the WebGPU op-test framework.
+
+HF Linear lowers to addmm (not aten.linear), so this is the dominant GEMM in the
+Florence-2 BART + DaViT graphs: out = beta*self + alpha*(mat1 @ mat2).
+"""
+
+import torch
+
+
+class AddmmModule(torch.nn.Module):
+    def __init__(self, n: int, beta: float = 1.0, alpha: float = 1.0):
+        super().__init__()
+        self.bias = torch.nn.Parameter(torch.linspace(-0.5, 0.5, n))
+        self.beta, self.alpha = beta, alpha
+
+    def forward(self, mat1: torch.Tensor, mat2: torch.Tensor) -> torch.Tensor:
+        return torch.addmm(self.bias, mat1, mat2, beta=self.beta, alpha=self.alpha)
+
+
+def _randn(shape) -> torch.Tensor:
+    g = torch.Generator().manual_seed(sum(int(x) for x in shape))
+    return torch.randn(*shape, generator=g) * 0.1