[ExecuTorch][WebGPU] Op-tests for q8ta_add

backends/webgpu/test/op_tests/cases.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -49,6 +49,7 @@
     DequantizeConstModule,
     QuantizeModule,
 )
+from executorch.backends.webgpu.test.ops.test_q8ta_add import Q8taAddModule
 from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule
 from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule
 from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule
@@ -761,6 +762,29 @@
     )
 
 
+# span the int8 sign-extend + requant-clamp edges (-128/127); numel % 4 == 0.
+_Q8_A = [-128, 127, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7]
+_Q8_B = [-128, 127, 0, 1, -1, 50, -50, 100, -100, 3, -3, 42, -42, 9, -9, 63]
+
+
+@register_op_test("q8ta_add")
+def _q8ta_add_suite() -> WebGPUTestSuite:
+    # int8 add (baked int8 consts), byte-exact vs CPU eager. `alpha` case pins the
+    # a + alpha*b term (the Vulkan glsl buffer path drops alpha).
+    def case(name, **kw):
+        return Case(name=name, construct={"a_vals": _Q8_A, "b_vals": _Q8_B, **kw})
+
+    return WebGPUTestSuite(
+        module_factory=lambda **kw: Q8taAddModule(**kw),
+        cases=[
+            case("basic"),
+            case("alpha", alpha=2.0),
+            case("nonzero_zp", a_zp=5, b_zp=-4, out_zp=7),
+        ],
+        golden_dtype="float32",
+    )
+
+
 @register_op_test("dequantize_per_tensor")
 def _dequant_const_suite() -> WebGPUTestSuite:
     # Independent GPU dequantize check: baked int8 const vs torch dequant (breaks

backends/webgpu/test/ops/test_q8ta_add.py

@@ -0,0 +1,70 @@
+# 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.
+
+"""`et_vk.q8ta_add` module + configs: int8 elementwise add.
+
+Both int8 operands are baked constants (only the scalar qparams vary), so the op
+is exercised alone with a byte-exact int8 golden vs the CPU eager op. The kernel
+implements `a + alpha*b` (the CPU reference semantics); an `alpha != 1` case pins
+the alpha term, which the Vulkan glsl buffer path drops.
+"""
+
+import unittest
+
+import torch
+
+import executorch.backends.vulkan.custom_ops_lib  # noqa: F401 registers et_vk ops
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+
+class Q8taAddModule(torch.nn.Module):
+    def __init__(
+        self,
+        a_vals,
+        b_vals,
+        a_scale=0.05,
+        a_zp=0,
+        b_scale=0.1,
+        b_zp=3,
+        out_scale=0.08,
+        out_zp=-2,
+        alpha=1.0,
+    ):
+        super().__init__()
+        self.register_buffer("a", torch.tensor(a_vals, dtype=torch.int8))
+        self.register_buffer("b", torch.tensor(b_vals, dtype=torch.int8))
+        self.a_scale, self.a_zp = a_scale, a_zp
+        self.b_scale, self.b_zp = b_scale, b_zp
+        self.out_scale, self.out_zp, self.alpha = out_scale, out_zp, alpha
+
+    def forward(self) -> torch.Tensor:
+        return torch.ops.et_vk.q8ta_add(
+            self.a,
+            self.b,
+            self.a_scale,
+            self.a_zp,
+            self.b_scale,
+            self.b_zp,
+            self.out_scale,
+            self.out_zp,
+            self.alpha,
+        )
+
+
+class Q8taAddTest(unittest.TestCase):
+    def test_delegates(self) -> None:
+        m = Q8taAddModule(list(range(-8, 8)), list(range(7, -9, -1)))
+        et = to_edge_transform_and_lower(
+            torch.export.export(m, ()), partitioner=[VulkanPartitioner()]
+        ).to_executorch()
+        delegate_ids = [
+            d.id
+            for plan in et.executorch_program.execution_plan
+            for d in plan.delegates
+        ]
+        self.assertIn("VulkanBackend", delegate_ids)