diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 16e331aa268..4bb392c0f02 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -527,3 +527,22 @@ def _embedding_suite() -> WebGPUTestSuite: 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, + ) diff --git a/backends/webgpu/test/ops/test_addmm.py b/backends/webgpu/test/ops/test_addmm.py new file mode 100644 index 00000000000..118ccaa1b92 --- /dev/null +++ b/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