From 2a0675a323ca92c78658a4c20a53faaf8cd033b2 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Sat, 4 Jul 2026 18:19:47 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/test/test_wgsl_codegen.py | 243 ++++++++++++++++++++++ 1 file changed, 243 insertions(+) diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 7d32a230459..392ddc174b5 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -11,6 +11,8 @@ import hashlib import importlib.util +import json +import re import tempfile import unittest from pathlib import Path @@ -20,6 +22,44 @@ g = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(g) +# gen_wgsl_headers.py and backends/vulkan/runtime/gen_vulkan_spv.py share the +# same $-block transpiler helpers; the test below keeps them in sync with that +# source of truth. Resolve the path relative to the repo root (both +# backends/vulkan and backends/webgpu exist in pytorch/executorch) and compare +# the bodies as TEXT -- gen_vulkan_spv.py cannot be imported (it top-level +# `import yaml`s, absent on the codegen runtime). +_REPO_ROOT = g.BACKEND_ROOT.parents[1] +_VULKAN_SPV = _REPO_ROOT / "backends" / "vulkan" / "runtime" / "gen_vulkan_spv.py" +_SHARED_TRANSPILER_FNS = ("extract_leading_whitespace", "escape", "preprocess") + + +def _function_source(text: str, name: str) -> str: + """Return a top-level function's source: its `def ` line through the + last line before the next column-0 construct (def-line to next dedent). + + Two-phase so a multi-line signature -- whose closing `) -> str:` sits at + column 0 -- is not mistaken for the next top-level construct. + """ + lines = text.splitlines() + start = next( + (i for i, ln in enumerate(lines) if re.match(rf"^def {re.escape(name)}\b", ln)), + None, + ) + if start is None: + raise AssertionError(f"def {name} not found") + # Advance past the (possibly multi-line) signature to the line ending in ':'. + sig = start + while not lines[sig].rstrip().endswith(":"): + sig += 1 + # The body ends at the next non-blank column-0 line. + end = len(lines) + for k in range(sig + 1, len(lines)): + head = lines[k][:1] + if head != "" and not head.isspace(): + end = k + break + return "\n".join(lines[start:end]).rstrip() + class WgslCodegenTest(unittest.TestCase): def test_symbol_base(self) -> None: @@ -189,5 +229,208 @@ def test_render_header_3d_emits_xyz(self) -> None: self.assertIn("inline constexpr uint32_t kFooWorkgroupSizeZ = 2;", h) +class WgslTemplateEngineTest(unittest.TestCase): + """Coverage for the $-block template engine + DTYPE/VEC variant matrix.""" + + # --- transpiler helpers stay in sync with their source --- + + @unittest.skipUnless( + _VULKAN_SPV.exists(), f"source of truth not present at {_VULKAN_SPV}" + ) + def test_transpiler_helpers_stay_in_sync(self) -> None: + # The shared $-block transpiler helpers must stay character-identical to + # their source of truth so they cannot silently drift. Read both files as + # TEXT (the source of truth cannot be imported -- it top-level + # `import yaml`s). + src_text = _VULKAN_SPV.read_text() + gen_text = _GEN.read_text() + for fn in _SHARED_TRANSPILER_FNS: + self.assertEqual( + _function_source(src_text, fn), + _function_source(gen_text, fn), + f"{fn} has drifted from its source of truth " + f"({_VULKAN_SPV}) -- re-sync the shared transpiler helpers", + ) + + # --- preprocess ------------------------------------------------------- + + def test_preprocess_if_else_selects_branch(self) -> None: + tmpl = 'fn main() {\n $if MODE == "a":\n let x = 1;\n $else:\n let x = 2;\n}\n' + self.assertEqual( + g.preprocess(tmpl, {"MODE": "a"}), "fn main() {\n let x = 1;\n}\n" + ) + self.assertEqual( + g.preprocess(tmpl, {"MODE": "b"}), "fn main() {\n let x = 2;\n}\n" + ) + + def test_preprocess_inline_substitution_uses_helper(self) -> None: + tmpl = "type: ${wgsl_buffer_type(DTYPE, VEC)};\n" + out = g.preprocess(tmpl, {**g.WGSL_HELPERS, "DTYPE": "float", "VEC": 4}) + self.assertEqual(out, "type: vec4;\n") + + def test_preprocess_guarded_body_indent_matches_control_column(self) -> None: + # $if authored at column 2 with its body one 2-space level deeper -> the + # guarded output line lands at column 2 (the control-line's column). + tmpl = "fn main() {\n $if VEC == 4:\n let a = 1;\n $else:\n let b = 2;\n}\n" + self.assertEqual( + g.preprocess(tmpl, {"VEC": 4}), "fn main() {\n let a = 1;\n}\n" + ) + self.assertEqual( + g.preprocess(tmpl, {"VEC": 1}), "fn main() {\n let b = 2;\n}\n" + ) + + def test_preprocess_enable_f16_only_for_half(self) -> None: + # DD-009: `enable f16;` is a literal line behind `$if DTYPE == "half":`, + # NOT an inline ${} (which would print a stray blank line for float and + # break byte-identity of the fp32 base). + tmpl = '$if DTYPE == "half":\n enable f16;\nfn main() {}\n' + self.assertEqual( + g.preprocess(tmpl, {"DTYPE": "half"}), "enable f16;\nfn main() {}\n" + ) + self.assertEqual(g.preprocess(tmpl, {"DTYPE": "float"}), "fn main() {}\n") + + # --- generate_variant_combinations ----------------------------------- + + def test_generate_variant_combinations_product(self) -> None: + iterated = { + "DTYPE": [{"VALUE": "float"}, {"VALUE": "half", "SUFFIX": "half"}], + "VEC": [{"VALUE": 1, "SUFFIX": ""}, {"VALUE": 4, "SUFFIX": "vec4"}], + } + combos = g.generate_variant_combinations(iterated) + self.assertEqual(len(combos), 4) + flat = [tuple((s[0], s[1], s[2]) for s in combo) for combo in combos] + self.assertIn((("DTYPE", "float", "float"), ("VEC", "", 1)), flat) + self.assertIn((("DTYPE", "half", "half"), ("VEC", "vec4", 4)), flat) + + def test_generate_variant_combinations_suffix_empty_suppresses(self) -> None: + combos = g.generate_variant_combinations({"VEC": [{"VALUE": 1, "SUFFIX": ""}]}) + self.assertEqual(combos, [(("VEC", "", 1),)]) + + def test_generate_variant_combinations_suffix_defaults_to_value(self) -> None: + # SUFFIX absent -> the suffix defaults to the VALUE (stringified in names). + combos = g.generate_variant_combinations({"VEC": [{"VALUE": 4}]}) + self.assertEqual(len(combos), 1) + ((name, suffix, value),) = combos[0] + self.assertEqual(name, "VEC") + self.assertEqual(value, 4) + self.assertEqual(str(suffix), "4") + + def test_generate_variant_combinations_excludes_param(self) -> None: + # A param already fixed by the variant is excluded from the forall product. + combos = g.generate_variant_combinations( + {"VEC": [{"VALUE": 1}, {"VALUE": 4}]}, {"VEC"} + ) + self.assertEqual(combos, [()]) + + # --- parse_template_spec --------------------------------------------- + + def _write_spec(self, tmp: str, name: str, spec_obj) -> Path: + p = Path(tmp) / f"{name}.json" + p.write_text(json.dumps(spec_obj)) + return p + + def test_parse_template_spec_minimal(self) -> None: + spec_obj = { + "op": { + "parameter_names_with_default_values": {"DTYPE": "float", "VEC": 1}, + "generate_variant_forall": { + "VEC": [ + {"VALUE": 1, "SUFFIX": ""}, + {"VALUE": 4, "SUFFIX": "vec4"}, + ] + }, + "shader_variants": [{"NAME": "op"}], + } + } + with tempfile.TemporaryDirectory() as tmp: + parsed = g.parse_template_spec(self._write_spec(tmp, "op", spec_obj)) + self.assertEqual(list(parsed.keys()), ["op"]) + v1, v4 = parsed["op"] + self.assertEqual((v1["NAME"], v1["VEC"], v1["DTYPE"]), ("op", 1, "float")) + self.assertEqual(v1["VARIANT_NAME"], "op") + self.assertEqual((v4["NAME"], v4["VEC"], v4["DTYPE"]), ("op_vec4", 4, "float")) + self.assertEqual(v4["VARIANT_NAME"], "op") + + def test_parse_template_spec_default_suffix_str_value_in_name(self) -> None: + # A forall value with no SUFFIX contributes str(VALUE) to the variant NAME. + spec_obj = { + "op": { + "parameter_names_with_default_values": {"VEC": 1}, + "generate_variant_forall": {"VEC": [{"VALUE": 4}]}, + "shader_variants": [{"NAME": "op"}], + } + } + with tempfile.TemporaryDirectory() as tmp: + parsed = g.parse_template_spec(self._write_spec(tmp, "op", spec_obj)) + self.assertEqual(parsed["op"][0]["NAME"], "op_4") + + def test_parse_template_spec_duplicate_key_raises(self) -> None: + # The dup-key object_pairs_hook rejects a repeated key anywhere in the spec. + dup = '{"op": {"NAME": 1, "NAME": 2}}' + with tempfile.TemporaryDirectory() as tmp: + p = Path(tmp) / "op.json" + p.write_text(dup) + with self.assertRaises(ValueError): + g.parse_template_spec(p) + + def test_headers_for_shader_top_level_key_must_match_stem(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + op_dir = Path(tmp) / "runtime/ops/op" + op_dir.mkdir(parents=True) + (op_dir / "op.wgsl").write_text("@workgroup_size(64)\nfn main(){}\n") + # top-level key "WRONG" != stem "op" -> must raise. + (op_dir / "op.json").write_text( + '{"WRONG": {"parameter_names_with_default_values": {},' + ' "shader_variants": [{"NAME": "op"}]}}' + ) + with self.assertRaises(ValueError): + list(g.headers_for_shader(op_dir / "op.wgsl")) + + def test_headers_for_shader_templating_without_sidecar_raises(self) -> None: + # A $if/${ shader with no sibling .json spec is a hard error. + with tempfile.TemporaryDirectory() as tmp: + op_dir = Path(tmp) / "runtime/ops/op" + op_dir.mkdir(parents=True) + (op_dir / "op.wgsl").write_text( + "$if VEC == 4:\n x\n@workgroup_size(64)\nfn main(){}\n" + ) + with self.assertRaises(ValueError): + list(g.headers_for_shader(op_dir / "op.wgsl")) + + # --- WGSL type-helpers ----------------------------------------------- + + def test_wgsl_scalar_type(self) -> None: + self.assertEqual(g.wgsl_scalar_type("half"), "f16") + self.assertEqual(g.wgsl_scalar_type("float"), "f32") + + def test_wgsl_buffer_type(self) -> None: + self.assertEqual(g.wgsl_buffer_type("float", 1), "f32") + self.assertEqual(g.wgsl_buffer_type("float", 4), "vec4") + self.assertEqual(g.wgsl_buffer_type("half", 4), "vec4") + + def test_wgsl_accum_type_is_f32(self) -> None: + self.assertEqual(g.wgsl_accum_type(), "f32") + + # --- byte-identity round-trip ---------------------------------------- + + def test_rms_norm_template_roundtrip_byte_identical(self) -> None: + # Expanding the committed rms_norm.wgsl template + embedding it must + # reproduce the committed headers exactly (the dedup proof point). + rms_dir = g.BACKEND_ROOT / "runtime/ops/rms_norm" + template = (rms_dir / "rms_norm.wgsl").read_text() + for name, vec, header_name in [ + ("rms_norm", 1, "rms_norm_wgsl.h"), + ("rms_norm_vec4", 4, "rms_norm_vec4_wgsl.h"), + ]: + expanded = g.preprocess( + template, {**g.WGSL_HELPERS, "DTYPE": "float", "VEC": vec} + ) + want = g.render_header(name, expanded, "rms_norm") + got = (rms_dir / header_name).read_text() + self.assertEqual( + got, want, f"{header_name} not reproduced from rms_norm.wgsl template" + ) + + if __name__ == "__main__": unittest.main()