From 8635a76267fa3b04ee0fb731b007d2dbc2b09140 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 28 Jul 2026 16:59:03 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../webgpu/runtime/WebGPUShaderRegistry.cpp | 4 +- .../runtime/ops/logical_and/LogicalAnd.cpp | 2 +- .../logical_and_wgsl.h | 2 +- .../logical_binary.wgsl} | 7 +- .../ops/logical_binary/logical_binary.yaml | 13 +++ .../logical_or_wgsl.h | 2 +- .../runtime/ops/logical_or/LogicalOr.cpp | 2 +- .../runtime/ops/logical_or/logical_or.wgsl | 25 ----- backends/webgpu/test/op_tests/cases.py | 91 +++++-------------- .../webgpu/test/op_tests/test_generator.py | 49 ++++++++++ backends/webgpu/test/ops/test_bitwise.py | 20 ++-- backends/webgpu/test/ops/test_logical_and.py | 51 ++++++----- backends/webgpu/test/ops/test_logical_or.py | 42 ++------- backends/webgpu/test/test_wgsl_codegen.py | 64 +++++++++++-- 14 files changed, 201 insertions(+), 173 deletions(-) rename backends/webgpu/runtime/ops/{logical_and => logical_binary}/logical_and_wgsl.h (96%) rename backends/webgpu/runtime/ops/{logical_and/logical_and.wgsl => logical_binary/logical_binary.wgsl} (73%) create mode 100644 backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml rename backends/webgpu/runtime/ops/{logical_or => logical_binary}/logical_or_wgsl.h (96%) delete mode 100644 backends/webgpu/runtime/ops/logical_or/logical_or.wgsl diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index b475a6cd999..4efd28dd05c 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -65,8 +65,8 @@ #include #include #include -#include -#include +#include +#include #include #include #include diff --git a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp index b6a7e12830f..358cb29fff0 100644 --- a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp +++ b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h b/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h similarity index 96% rename from backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h rename to backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h index 6a21a77a687..3c1f861f119 100644 --- a/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h +++ b/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from logical_and.wgsl - DO NOT EDIT. +// @generated from logical_binary.wgsl - DO NOT EDIT. // wgsl-sha256: cf7c1d1dbba94e429120796c9c25a6717786cca03c08f3bd1e291d5627089c20 inline constexpr const char* kLogicalAndWGSL = R"( @group(0) @binding(0) var t_out: array; diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl b/backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl similarity index 73% rename from backends/webgpu/runtime/ops/logical_and/logical_and.wgsl rename to backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl index 9acb583f51c..fb5310cfde9 100644 --- a/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl +++ b/backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl @@ -16,10 +16,13 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. + $if OP == "&": + // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. + $else: + // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. let w = gid.x + gid.y * (num_workgroups.x * wg_size); if (w >= params.num_words) { return; } - t_out[w] = t_a[w] & t_b[w]; + t_out[w] = t_a[w] ${OP} t_b[w]; } diff --git a/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml b/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml new file mode 100644 index 00000000000..65df4465b04 --- /dev/null +++ b/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml @@ -0,0 +1,13 @@ +# 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. + +logical_binary: + parameter_names_with_default_values: + OP: "&" + shader_variants: + - NAME: logical_and + - NAME: logical_or + OP: "|" diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h b/backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h similarity index 96% rename from backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h rename to backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h index d64898cb523..e61317d2ff1 100644 --- a/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h +++ b/backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from logical_or.wgsl - DO NOT EDIT. +// @generated from logical_binary.wgsl - DO NOT EDIT. // wgsl-sha256: 4ad19ee04e2c7b396b4669cf44f95133d658c3ec2e6f37d7b271bedc0e582ecf inline constexpr const char* kLogicalOrWGSL = R"( @group(0) @binding(0) var t_out: array; diff --git a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp index d26f6b486d2..d750af11227 100644 --- a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp +++ b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl b/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl deleted file mode 100644 index d7e6176ba32..00000000000 --- a/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl +++ /dev/null @@ -1,25 +0,0 @@ -@group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_a: array; -@group(0) @binding(2) var t_b: array; - -struct Params { - num_words: u32, - pad0: u32, - pad1: u32, - pad2: u32, -} -@group(0) @binding(3) var params: Params; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. - let w = gid.x + gid.y * (num_workgroups.x * wg_size); - if (w >= params.num_words) { - return; - } - t_out[w] = t_a[w] | t_b[w]; -} diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index aa8f000bceb..f429de949ac 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -38,10 +38,10 @@ ) from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule from executorch.backends.webgpu.test.ops.test_bitwise import ( + BITWISE_NOT_SHAPES, BitwiseAndModule, BitwiseNotModule, bw_gen_a, - bw_gen_b, ) from executorch.backends.webgpu.test.ops.test_cat import ( CatModule, @@ -71,14 +71,13 @@ make_qcs4w_linear_module, ) from executorch.backends.webgpu.test.ops.test_logical_and import ( - la_gen_a, - la_gen_b, + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, LogicalAndModule, ) from executorch.backends.webgpu.test.ops.test_logical_or import ( BitwiseOrModule, - lo_gen_a, - lo_gen_b, LogicalOrModule, ) from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule @@ -334,49 +333,33 @@ def _ge_suite() -> WebGPUTestSuite: return _compare_suite("ge") -@register_op_test("logical_and") -def _logical_and_suite() -> WebGPUTestSuite: - # out = (a>0) && (b>0): two bool masks derived on-GPU from float inputs via - # gt.Tensor (baked zeros), AND'd -> bool. Distinct a/b seeds so the masks - # differ (AND ~25% True, a real mix an OR mutant fails); all shapes numel % - # 4 == 0 (bool packs 4/word). float32 oracle (byte-exact bool golden). +def _logical_binary_suite(module_factory) -> WebGPUTestSuite: + # Every packed u32 receives all four Boolean pairs in byte order. def case(name, shape): return Case( name=name, construct={"shape": shape}, inputs=( - InputSpec(shape=shape, gen=la_gen_a), - InputSpec(shape=shape, gen=la_gen_b), + InputSpec(shape=shape, gen=logical_binary_gen_a), + InputSpec(shape=shape, gen=logical_binary_gen_b), ), ) return WebGPUTestSuite( - module_factory=lambda shape: LogicalAndModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + module_factory=module_factory, + cases=[case(name, shape) for name, shape in LOGICAL_BINARY_CASES], golden_dtype="float32", ) +@register_op_test("logical_and") +def _logical_and_suite() -> WebGPUTestSuite: + return _logical_binary_suite(lambda shape: LogicalAndModule(shape)) + + @register_op_test("bitwise_and") def _bitwise_and_suite() -> WebGPUTestSuite: - # bool bitwise AND == logical_and for canonical 0/1 (shares the handler). - # Two masks derived on-GPU from float inputs via gt.Tensor (baked zeros), - # distinct a/b seeds (AND ~25% True); all shapes numel % 4 == 0. - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=bw_gen_a), - InputSpec(shape=shape, gen=bw_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: BitwiseAndModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: BitwiseAndModule(shape)) @register_op_test("bitwise_not") @@ -392,52 +375,22 @@ def case(name, shape): return WebGPUTestSuite( module_factory=lambda shape: BitwiseNotModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + cases=[ + case(name, shape) + for name, shape in zip(("2d", "3d", "sq"), BITWISE_NOT_SHAPES) + ], golden_dtype="float32", ) @register_op_test("logical_or") def _logical_or_suite() -> WebGPUTestSuite: - # out = (a>0) || (b>0): two bool masks derived on-GPU from float inputs via - # gt.Tensor (baked zeros), OR'd -> bool. Distinct a/b seeds (~50% each, - # independent -> OR ~75% True, a real mix an AND mutant fails); all shapes - # numel % 4 == 0. float32 oracle (byte-exact bool golden). - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=lo_gen_a), - InputSpec(shape=shape, gen=lo_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: LogicalOrModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: LogicalOrModule(shape)) @register_op_test("bitwise_or") def _bitwise_or_suite() -> WebGPUTestSuite: - # bool bitwise OR == logical_or for canonical 0/1 (shares the handler). - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=lo_gen_a), - InputSpec(shape=shape, gen=lo_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: BitwiseOrModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: BitwiseOrModule(shape)) @register_op_test("pow") diff --git a/backends/webgpu/test/op_tests/test_generator.py b/backends/webgpu/test/op_tests/test_generator.py index fabf79171c0..eaf9b72d127 100644 --- a/backends/webgpu/test/op_tests/test_generator.py +++ b/backends/webgpu/test/op_tests/test_generator.py @@ -11,6 +11,11 @@ from executorch.backends.webgpu.test.op_tests import generate_op_tests as g from executorch.backends.webgpu.test.op_tests.test_suite import op_test_registry +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) def _add_regular_case(): @@ -112,3 +117,47 @@ def test_manifest_schema_roundtrip(tmp_path): gd = e["golden"] assert {"path", "shape", "dtype", "output_index"} <= set(gd) assert gd["output_index"] == 0 + + +def test_logical_binary_case_contract(): + expected_cases = ( + ("2d", (4, 8)), + ("3d", (2, 3, 8)), + ("sq", (16, 16)), + ("words63", (252,)), + ("words64", (256,)), + ("words65", (260,)), + ) + assert LOGICAL_BINARY_CASES == expected_cases + assert logical_binary_gen_a((8,)).tolist() == [ + -1.0, + -1.0, + 1.0, + 1.0, + -1.0, + -1.0, + 1.0, + 1.0, + ] + assert logical_binary_gen_b((8,)).tolist() == [ + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + ] + + for op in ("logical_and", "bitwise_and", "logical_or", "bitwise_or"): + suite = op_test_registry[op] + assert tuple((case.name, case.construct["shape"]) for case in suite.cases) == ( + expected_cases + ) + for case in suite.cases: + assert case.required is True + assert case.heavy is False + assert len(case.inputs) == 2 + assert case.inputs[0].gen is logical_binary_gen_a + assert case.inputs[1].gen is logical_binary_gen_b diff --git a/backends/webgpu/test/ops/test_bitwise.py b/backends/webgpu/test/ops/test_bitwise.py index b61a37d9990..a8fee7a0095 100644 --- a/backends/webgpu/test/ops/test_bitwise.py +++ b/backends/webgpu/test/ops/test_bitwise.py @@ -19,6 +19,11 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) from executorch.exir import to_edge_transform_and_lower @@ -50,11 +55,10 @@ def g(shape): bw_gen_a = _bw_gen(0) -bw_gen_b = _bw_gen(1) # All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] +BITWISE_NOT_SHAPES = ((4, 8), (2, 3, 8), (16, 16)) class BitwiseTest(unittest.TestCase): @@ -75,13 +79,17 @@ def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ) def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = bw_gen_a(shape) - b = bw_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(op="bitwise_and", case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) self._assert_delegates( BitwiseAndModule(shape), (a, b), "bitwise_and", shape ) + + for shape in BITWISE_NOT_SHAPES: + with self.subTest(op="bitwise_not", shape=shape): + a = bw_gen_a(shape) self._assert_delegates( BitwiseNotModule(shape), (a,), "bitwise_not", shape ) diff --git a/backends/webgpu/test/ops/test_logical_and.py b/backends/webgpu/test/ops/test_logical_and.py index 4c3939002d0..9cf474fd8de 100644 --- a/backends/webgpu/test/ops/test_logical_and.py +++ b/backends/webgpu/test/ops/test_logical_and.py @@ -4,17 +4,9 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.logical_and.default` module + configs for the WebGPU op-test framework. - -`LogicalAndModule` derives its two bool operands on-GPU from float inputs -(`a > 0`, `b > 0` via the delegated `gt.Tensor` against a baked zero buffer), so -the only runtime inputs are the two float tensors (the op-test framework is -float-input-only). `a`/`b` use distinct seeds so the two bool masks differ (each -~50% True, independent -> AND ~25% True), a real mix that a wrong op (e.g. OR) -would fail. Output is bool (byte-exact golden). `LogicalAndTest` is the -export-delegation smoke test. -""" +"""Delegation coverage for logical AND with packed truth-table inputs.""" +import math import unittest import torch @@ -32,29 +24,40 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.logical_and(a > self.z, b > self.z) -def _la_gen(seed): - # Distinct per-input seed so the two derived bool masks differ. - def g(shape): - gen = torch.Generator().manual_seed(seed) - return torch.randn(*shape, generator=gen, dtype=torch.float32) +LOGICAL_BINARY_CASES = ( + ("2d", (4, 8)), + ("3d", (2, 3, 8)), + ("sq", (16, 16)), + ("words63", (252,)), + ("words64", (256,)), + ("words65", (260,)), +) - return g +def _logical_binary_gen(pattern): + def generate(shape): + numel = math.prod(shape) + if numel == 0 or numel % len(pattern) != 0: + raise ValueError("logical-binary test shapes must have numel % 4 == 0") + return ( + torch.tensor(pattern, dtype=torch.float32) + .repeat(numel // len(pattern)) + .reshape(shape) + ) -la_gen_a = _la_gen(0) -la_gen_b = _la_gen(1) + return generate -# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] +logical_binary_gen_a = _logical_binary_gen((-1.0, -1.0, 1.0, 1.0)) +logical_binary_gen_b = _logical_binary_gen((-1.0, 1.0, -1.0, 1.0)) class LogicalAndTest(unittest.TestCase): def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = la_gen_a(shape) - b = la_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) ep = torch.export.export(LogicalAndModule(shape).eval(), (a, b)) edge = to_edge_transform_and_lower( ep, partitioner=[VulkanPartitioner()] diff --git a/backends/webgpu/test/ops/test_logical_or.py b/backends/webgpu/test/ops/test_logical_or.py index d71c1493ea9..420fee26336 100644 --- a/backends/webgpu/test/ops/test_logical_or.py +++ b/backends/webgpu/test/ops/test_logical_or.py @@ -4,23 +4,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.logical_or.default` / `aten.bitwise_or.Tensor` (bool) modules + configs. - -Mirrors the logical_and/bitwise_and tests: the modules derive their two bool -operands on-GPU from float inputs (`a > 0`, `b > 0` via the delegated `gt.Tensor` -against a baked zero buffer), so the only runtime inputs are the two float -tensors (the op-test framework is float-input-only). `a`/`b` use distinct seeds -so the two bool masks differ (each ~50% True, independent -> OR ~75% True), a -real mix that a wrong op (e.g. AND) would fail. `bitwise_or` on bool is identical -to `logical_or` (shares the handler). Output is bool (byte-exact golden). -`LogicalOrTest` is the export-delegation smoke test. -""" +"""Delegation coverage for logical and bitwise OR truth-table inputs.""" import unittest import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) from executorch.exir import to_edge_transform_and_lower @@ -42,23 +37,6 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.bitwise_or(a > self.z, b > self.z) -def _lo_gen(seed): - # Distinct per-input seed so the two derived bool masks differ. - def g(shape): - gen = torch.Generator().manual_seed(seed) - return torch.randn(*shape, generator=gen, dtype=torch.float32) - - return g - - -lo_gen_a = _lo_gen(0) -lo_gen_b = _lo_gen(1) - - -# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] - - class LogicalOrTest(unittest.TestCase): def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ep = torch.export.export(mod.eval(), inputs) @@ -77,10 +55,10 @@ def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ) def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = lo_gen_a(shape) - b = lo_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) self._assert_delegates( LogicalOrModule(shape), (a, b), "logical_or", shape ) diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 7305142c4d1..49be223ac46 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -230,7 +230,11 @@ def test_generated_output_manifest_digest(self) -> None: self.assertEqual(len(outputs), 134) self.assertEqual( digest.hexdigest(), - "a3f6324b224c049f239beea3c38e1dd532e364f8d44c595eb5fcc84dba4a1a83", + "6c2b6fa24aca2de0ae395979d15015155afe81b64d51b3fee81368da315ceaf6", + ) + self.assertEqual( + hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), + "e82cab9afb45a880e3f33da73ff0d000bee8092eccca3ccf5819701bb433d6d7", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: @@ -967,10 +971,6 @@ def test_to_copy_convert_template_roundtrip_byte_identical(self) -> None: entries["to_copy_int_to_float"].include, "runtime/ops/to_copy/to_copy_int_to_float_wgsl.h", ) - self.assertEqual( - hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", - ) def test_extrema_template_roundtrip_byte_identical(self) -> None: extrema_dir = g.BACKEND_ROOT / "runtime/ops/extrema" @@ -1021,10 +1021,56 @@ def test_extrema_template_roundtrip_byte_identical(self) -> None: hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash ) - self.assertEqual( - hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", - ) + def test_logical_binary_template_roundtrip_byte_identical(self) -> None: + logical_dir = g.BACKEND_ROOT / "runtime/ops/logical_binary" + template_path = logical_dir / "logical_binary.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "logical_and": ( + "&", + "cf7c1d1dbba94e429120796c9c25a6717786cca03c08f3bd1e291d5627089c20", + ), + "logical_or": ( + "|", + "4ad19ee04e2c7b396b4669cf44f95133d658c3ec2e6f37d7b271bedc0e582ecf", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + + for name, (op, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["OP"], op) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header_path = logical_dir / f"{name}_wgsl.h" + header = header_path.read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (64, 1, 1)) + + entries = {entry.name: entry for entry in g.registry_entries()} + for name in expected: + self.assertEqual( + entries[name].include, + f"runtime/ops/logical_binary/{name}_wgsl.h", + ) + self.assertEqual(entries[name].symbol, g.symbol_base(name)) + + handler_hashes = { + "logical_and": "eb85a8f97ee7640298a661da49feb08aa79b8c24d3d4458b71d24d3f01bc388d", + "logical_or": "bda18617f7077fee5a812c21cdc495c89542a1688f7e1ef6739ed01da343a66b", + } + for name, expected_hash in handler_hashes.items(): + handler = ( + g.BACKEND_ROOT / f"runtime/ops/{name}/Logical{name[8:].title()}.cpp" + ) + self.assertEqual( + hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash + ) def test_rms_norm_half_variant_is_type_correct(self) -> None: # A DTYPE=half expansion must emit compilable WGSL: `enable f16;`, an f32