[ExecuTorch][WebGPU] Tests for the matmul WebGPU ops

Differential Revision: D111755120

Pull Request resolved: #20918

Author: JCNTH

backends/webgpu/test/ops/test_bmm.py

@@ -0,0 +1,104 @@
+# 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.bmm.default` (fp32 batched GEMM) export + golden for the WebGPU backend.
+
+Exports single-op batched-matmul graphs through VulkanPartitioner and writes a
+torch-computed fp64 golden (the native binary has no ATen) + the raw fp32 inputs
+the native test loads and compares. Configs span the vec4 fast path (K%4==0 and
+N%4==0, wider 16B loads) and the scalar tiled path (K or N not a multiple of 4,
+which also exercises the bounds guards), across batch sizes 1, 2, and 4.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (shape_a [B, M, K], shape_b [B, K, N]). Output is [B, M, N].
+CONFIGS = {
+    "square": ((2, 64, 64), (2, 64, 64)),  # vec4 path, multi-tile, B=2
+    "tall": ((4, 128, 16), (4, 16, 24)),  # vec4 path, tall M, B=4
+    "scalar": ((2, 7, 5), (2, 5, 3)),  # scalar tiled path (K,N %4!=0)
+    "batch1": ((1, 33, 17), (1, 17, 5)),  # B=1, scalar path, bounds guards
+}
+
+
+class BmmModule(torch.nn.Module):
+    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return torch.bmm(a, b)
+
+
+def _det_inputs(shape_a, shape_b):
+    """Deterministic fp32 inputs (fixed seed) for a config."""
+    g = torch.Generator().manual_seed(0)
+    a = torch.randn(*shape_a, generator=g, dtype=torch.float32)
+    b = torch.randn(*shape_b, generator=g, dtype=torch.float32)
+    return a, b
+
+
+def _fp64_golden(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """fp64 truth: bmm(a, b) accumulated in double, cast back to fp32."""
+    return torch.bmm(a.double(), b.double()).to(torch.float32)
+
+
+def _export(a: torch.Tensor, b: torch.Tensor):
+    ep = torch.export.export(BmmModule().eval(), (a, b))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegated(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class BmmTest(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (sa, sb) in CONFIGS.items():
+            with self.subTest(name=name):
+                a, b = _det_inputs(sa, sb)
+                et = _export(a, b)
+                self.assertTrue(
+                    _delegated(et), f"Expected a VulkanBackend delegate (bmm {name})"
+                )
+
+    def test_golden_matches_fp64(self) -> None:
+        for name, (sa, sb) in CONFIGS.items():
+            with self.subTest(name=name):
+                a, b = _det_inputs(sa, sb)
+                torch.testing.assert_close(
+                    BmmModule()(a, b), _fp64_golden(a, b), atol=1e-4, rtol=1e-3
+                )
+
+
+def export_bmm_model(pte_path: str, golden_path: str, input_path: str) -> None:
+    """Write a bmm .pte + torch fp64 golden (raw LE fp32) + raw LE fp32 inputs (a then b)."""
+    a, b = _det_inputs(*CONFIGS["square"])
+    et = _export(a, b)
+    golden = _fp64_golden(a, b).numpy().astype("<f4")
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    golden.tofile(golden_path)
+    with open(input_path, "wb") as f:
+        f.write(a.numpy().astype("<f4").tobytes())
+        f.write(b.numpy().astype("<f4").tobytes())
+    print(
+        f"Exported {pte_path}; golden {golden_path} ({golden.size} floats); "
+        f"inputs {input_path} ({a.numel() + b.numel()} floats)"
+    )
+
+
+if __name__ == "__main__":
+    unittest.main()

backends/webgpu/test/ops/test_linear.py

@@ -0,0 +1,116 @@
+# 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, no bias) export + golden for the WebGPU backend.
+
+Exports single-op linear graphs through VulkanPartitioner and locks the
+delegation contract + an fp64 torch golden. The handler computes
+out[m,n] = sum_k x[m,k] * w[n,k] (weight is [N,K]); it picks a vec4-over-K
+kernel when K % 4 == 0 and a scalar tiled kernel otherwise, so the configs span
+both branches (plus a K==1 / non-multiple shape that exercises the bounds guard).
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (M, K, N). weight is [N, K]; output is [M, N].
+CONFIGS = {
+    "square": (64, 64, 64),  # K % 4 == 0 -> vec4-over-K kernel
+    "tall_vec4": (32, 128, 16),  # K % 4 == 0 -> vec4 path, skinny N
+    "scalar_k": (8, 7, 5),  # K % 4 != 0 -> scalar tiled path + bounds guard
+    "k1": (4, 1, 8),  # K == 1 rank-1, scalar path
+}
+
+
+class LinearModule(torch.nn.Module):
+    def __init__(self, weight: torch.Tensor) -> None:
+        super().__init__()
+        n, k = weight.shape
+        self.linear = torch.nn.Linear(k, n, bias=False)
+        with torch.no_grad():
+            self.linear.weight.copy_(weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear(x)
+
+
+def _det_inputs(m: int, k: int, n: int):
+    """Deterministic fp32 input [M,K] + weight [N,K] (fixed seed)."""
+    g = torch.Generator().manual_seed(0)
+    x = torch.randn(m, k, generator=g, dtype=torch.float32)
+    w = torch.randn(n, k, generator=g, dtype=torch.float32)
+    return x, w
+
+
+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()
+
+
+def _delegates(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestLinear(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (m, k, n) in CONFIGS.items():
+            with self.subTest(name=name):
+                x, w = _det_inputs(m, k, n)
+                et = _export(LinearModule(w).eval(), x)
+                self.assertTrue(
+                    _delegates(et), f"Expected a VulkanBackend delegate (linear {name})"
+                )
+
+    def test_golden_matches_fp64(self) -> None:
+        # Golden must match the fp64 F.linear truth (weight [N,K], no bias).
+        for name, (m, k, n) in CONFIGS.items():
+            with self.subTest(name=name):
+                x, w = _det_inputs(m, k, n)
+                got = LinearModule(w).eval()(x)
+                golden = torch.nn.functional.linear(x.double(), w.double())
+                torch.testing.assert_close(
+                    got, golden.to(torch.float32), atol=1e-3, rtol=1e-3
+                )
+
+
+def export_linear_model(
+    pte_path: str, golden_path: str, input_path: str, config: str = "square"
+) -> None:
+    """Write a linear .pte + torch fp64 golden (raw LE fp32) + raw LE fp32 input."""
+    m_dim, k, n = CONFIGS[config]
+    x, w = _det_inputs(m_dim, k, n)
+    model = LinearModule(w).eval()
+    golden = (
+        torch.nn.functional.linear(x.double(), w.double())
+        .to(torch.float32)
+        .numpy()
+        .astype("<f4")
+    )
+    et = _export(model, x)
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    golden.tofile(golden_path)
+    x.numpy().astype("<f4").tofile(input_path)
+    print(
+        f"Exported {pte_path}; golden {golden_path} ({golden.size} floats); "
+        f"input {input_path} ({x.numel()} floats)"
+    )
+
+
+if __name__ == "__main__":
+    unittest.main()

backends/webgpu/test/ops/test_matmul.py

@@ -0,0 +1,67 @@
+# 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.mm.default` (fp32 GEMM) module + configs for the WebGPU op-test framework.
+
+`MatmulModule` + `CONFIGS` are imported by `cases.py` to drive the declarative op-test
+suite (export via VulkanPartitioner + fp64 torch golden, run on Dawn). `MatmulTest` is
+the export-delegation smoke test. Configs span a square matmul, a tall/skinny shape,
+K==1 (rank-1 update), and non-power-of-two dims that exercise the bounds guard.
+"""
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (shape_a [M, K], shape_b [K, N]). Output is [M, N].
+CONFIGS = {
+    "square": ((64, 64), (64, 64)),
+    "tall": ((128, 32), (32, 16)),
+    "k1": ((8, 1), (1, 8)),
+    "nonmultiple": ((7, 5), (5, 3)),
+}
+
+
+class MatmulModule(torch.nn.Module):
+    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return torch.mm(a, b)
+
+
+def _det_inputs(shape_a, shape_b):
+    """Deterministic fp32 inputs (fixed seed) for a config."""
+    g = torch.Generator().manual_seed(0)
+    a = torch.randn(*shape_a, generator=g, dtype=torch.float32)
+    b = torch.randn(*shape_b, generator=g, dtype=torch.float32)
+    return a, b
+
+
+def _export(a: torch.Tensor, b: torch.Tensor):
+    ep = torch.export.export(MatmulModule().eval(), (a, b))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegated(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class MatmulTest(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (sa, sb) in CONFIGS.items():
+            with self.subTest(name=name):
+                a, b = _det_inputs(sa, sb)
+                et = _export(a, b)
+                self.assertTrue(
+                    _delegated(et), f"Expected a VulkanBackend delegate (mm {name})"
+                )