diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 6fa7125b68c..3ac13727afd 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -362,3 +362,50 @@ def _layer_norm_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) +from executorch.backends.webgpu.test.ops.test_linear_fp32 import ( + _ramp as _lin_ramp, + make_linear, +) + + +@register_op_test("linear_fp32") +def _linear_fp32_suite() -> WebGPUTestSuite: + # fp32 linear (BART + DaViT projections); bias + no-bias, and shapes whose + # M*N exceeds the 65535 1D ceiling to exercise the 2D-dispatch spill. + return WebGPUTestSuite( + module_factory=make_linear, + cases=[ + Case( + name="bias_mat", + construct={"in_features": 64, "out_features": 32}, + inputs=((4, 64),), + ), + Case( + name="no_bias", + construct={"in_features": 64, "out_features": 32, "bias": False}, + inputs=((4, 64),), + ), + Case( + name="rank3", + construct={"in_features": 768, "out_features": 768}, + inputs=(InputSpec(shape=(1, 16, 768), gen=_lin_ramp),), + ), + Case( + name="tall_m", + construct={"in_features": 128, "out_features": 64}, + inputs=((256, 128),), + ), + Case( + name="bart_proj", + construct={"in_features": 1024, "out_features": 1024}, + inputs=(InputSpec(shape=(1, 8, 1024), gen=_lin_ramp),), + ), + Case( + name="odd_k", + construct={"in_features": 63, "out_features": 32}, + inputs=((4, 63),), + ), + ], + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_linear_fp32.py b/backends/webgpu/test/ops/test_linear_fp32.py new file mode 100644 index 00000000000..c635b4f9a3e --- /dev/null +++ b/backends/webgpu/test/ops/test_linear_fp32.py @@ -0,0 +1,77 @@ +# 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.linear.default` (fp32) module + inputs for the WebGPU op-test framework. + +`make_linear` and `_ramp` are imported by `cases.py` to drive the declarative +op-test suite. `LinearFp32Test` is the export-delegation smoke test. fp32 linear +is the projection used throughout BART + the DaViT vision encoder (Florence-2); +both the bias and no-bias paths are covered. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +class LinearFp32Module(torch.nn.Module): + """fp32 linear; lowers to aten.linear.default with a prepacked [N, K] weight.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = True): + super().__init__() + self.fc = torch.nn.Linear(in_features, out_features, bias=bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc(x) + + +def make_linear( + in_features: int, out_features: int, bias: bool = True +) -> torch.nn.Module: + """Factory with deterministic weight (+ bias): a normalized ramp per row.""" + m = LinearFp32Module(in_features, out_features, bias=bias) + with torch.no_grad(): + w = torch.linspace( + -1.0, 1.0, out_features * in_features, dtype=torch.float32 + ).reshape(out_features, in_features) + m.fc.weight.copy_(w / in_features) + if bias: + m.fc.bias.copy_( + torch.linspace(-0.5, 0.5, out_features, dtype=torch.float32) + ) + return m + + +def _ramp(shape) -> torch.Tensor: + """Deterministic linear ramp in [-1, 1] reshaped to `shape`.""" + 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 LinearFp32Test(unittest.TestCase): + def test_export_delegates(self) -> None: + for bias in (True, False): + et = _export(make_linear(64, 32, bias=bias).eval(), _ramp((4, 64))) + 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 (linear bias={bias})" + )