diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl b/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl index 4bd5618596f..1234174981e 100644 --- a/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl +++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl @@ -1,6 +1,8 @@ -@group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_in: array; -@group(0) @binding(2) var t_weight: array; +$if DTYPE == "half": + enable f16; +@group(0) @binding(0) var t_out: array<${buffer_gvec_type(DTYPE, VEC)}>; +@group(0) @binding(1) var t_in: array<${buffer_gvec_type(DTYPE, VEC)}>; +@group(0) @binding(2) var t_weight: array<${buffer_gvec_type(DTYPE, VEC)}>; struct Params { num_rows: u32, @@ -12,7 +14,7 @@ struct Params { const WG_SIZE: u32 = 64u; -var shared_sum: array; +var shared_sum: array<${accum_scalar_type(DTYPE)}, WG_SIZE>; fn reduce_shared(worker_id: u32) { workgroupBarrier(); @@ -29,6 +31,10 @@ fn reduce_shared(worker_id: u32) { } } +$if VEC == 4: + // vec4 variant of rms_norm: each lane strides by WG_SIZE over rw4 = row_width/4 + // texels and accumulates dot(v, v). row_width is the ELEMENT count, so mean_sq + // divides by it (not rw4). The host selects this only when row_width % 4 == 0. @compute @workgroup_size(64, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, @@ -40,18 +46,39 @@ fn main( return; } - let base = row_idx * params.row_width; + $if VEC == 4: + let rw4 = params.row_width / 4u; + let base4 = row_idx * rw4; + $else: + let base = row_idx * params.row_width; - var local_sq_sum: f32 = 0.0; - var x: u32 = worker_id; - loop { - if (x >= params.row_width) { - break; + var local_sq_sum: ${accum_scalar_type(DTYPE)} = 0.0; + $if VEC == 4: + var x4: u32 = worker_id; + loop { + if (x4 >= rw4) { + break; + } + let v = t_in[base4 + x4]; + $if DTYPE == "half": + local_sq_sum = local_sq_sum + dot(vec4(v), vec4(v)); + $else: + local_sq_sum = local_sq_sum + dot(v, v); + x4 = x4 + WG_SIZE; + } + $else: + var x: u32 = worker_id; + loop { + if (x >= params.row_width) { + break; + } + let v = t_in[base + x]; + $if DTYPE == "half": + local_sq_sum = local_sq_sum + f32(v) * f32(v); + $else: + local_sq_sum = local_sq_sum + v * v; + x = x + WG_SIZE; } - let v = t_in[base + x]; - local_sq_sum = local_sq_sum + v * v; - x = x + WG_SIZE; - } shared_sum[worker_id] = local_sq_sum; reduce_shared(worker_id); @@ -59,14 +86,30 @@ fn main( let mean_sq = shared_sum[0] / f32(params.row_width); let rstd = inverseSqrt(mean_sq + params.epsilon); - x = worker_id; - loop { - if (x >= params.row_width) { - break; + $if VEC == 4: + x4 = worker_id; + loop { + if (x4 >= rw4) { + break; + } + $if DTYPE == "half": + t_out[base4 + x4] = vec4(vec4(t_in[base4 + x4]) * rstd * vec4(t_weight[x4])); + $else: + t_out[base4 + x4] = t_in[base4 + x4] * rstd * t_weight[x4]; + x4 = x4 + WG_SIZE; + } + $else: + x = worker_id; + loop { + if (x >= params.row_width) { + break; + } + let v = t_in[base + x]; + let w = t_weight[x]; + $if DTYPE == "half": + t_out[base + x] = f16(f32(v) * rstd * f32(w)); + $else: + t_out[base + x] = v * rstd * w; + x = x + WG_SIZE; } - let v = t_in[base + x]; - let w = t_weight[x]; - t_out[base + x] = v * rstd * w; - x = x + WG_SIZE; - } } diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm.yaml b/backends/webgpu/runtime/ops/rms_norm/rms_norm.yaml new file mode 100644 index 00000000000..6bfeeb012c3 --- /dev/null +++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm.yaml @@ -0,0 +1,12 @@ +rms_norm: + parameter_names_with_default_values: + DTYPE: float + VEC: 1 + generate_variant_forall: + VEC: + - VALUE: 1 + SUFFIX: "" + - VALUE: 4 + SUFFIX: vec4 + shader_variants: + - NAME: rms_norm diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4.wgsl b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4.wgsl deleted file mode 100644 index c2f731e5f60..00000000000 --- a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4.wgsl +++ /dev/null @@ -1,74 +0,0 @@ -@group(0) @binding(0) var t_out: array>; -@group(0) @binding(1) var t_in: array>; -@group(0) @binding(2) var t_weight: array>; - -struct Params { - num_rows: u32, - row_width: u32, - epsilon: f32, - _pad: u32, -} -@group(0) @binding(3) var params: Params; - -const WG_SIZE: u32 = 64u; - -var shared_sum: array; - -fn reduce_shared(worker_id: u32) { - workgroupBarrier(); - var stride: u32 = WG_SIZE / 2u; - loop { - if (stride == 0u) { - break; - } - if (worker_id < stride) { - shared_sum[worker_id] = shared_sum[worker_id] + shared_sum[worker_id + stride]; - } - workgroupBarrier(); - stride = stride >> 1u; - } -} - -// vec4 variant of rms_norm: each lane strides by WG_SIZE over rw4 = row_width/4 -// texels and accumulates dot(v, v). row_width is the ELEMENT count, so mean_sq -// divides by it (not rw4). The host selects this only when row_width % 4 == 0. -@compute @workgroup_size(64, 1, 1) -fn main( - @builtin(workgroup_id) wid: vec3, - @builtin(local_invocation_id) lid: vec3) { - let row_idx = wid.x; - let worker_id = lid.x; - - if (row_idx >= params.num_rows) { - return; - } - - let rw4 = params.row_width / 4u; - let base4 = row_idx * rw4; - - var local_sq_sum: f32 = 0.0; - var x4: u32 = worker_id; - loop { - if (x4 >= rw4) { - break; - } - let v = t_in[base4 + x4]; - local_sq_sum = local_sq_sum + dot(v, v); - x4 = x4 + WG_SIZE; - } - - shared_sum[worker_id] = local_sq_sum; - reduce_shared(worker_id); - - let mean_sq = shared_sum[0] / f32(params.row_width); - let rstd = inverseSqrt(mean_sq + params.epsilon); - - x4 = worker_id; - loop { - if (x4 >= rw4) { - break; - } - t_out[base4 + x4] = t_in[base4 + x4] * rstd * t_weight[x4]; - x4 = x4 + WG_SIZE; - } -} diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h index 633bf3adfc0..02213e3d7b0 100644 --- a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h +++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from rms_norm_vec4.wgsl - DO NOT EDIT. +// @generated from rms_norm.wgsl - DO NOT EDIT. // wgsl-sha256: 4c0ba56708bf125a7ec6ea3c51d1288e05ac00a8e2cfa10e38e9a208e230b8df inline constexpr const char* kRmsNormVec4WGSL = R"( @group(0) @binding(0) var t_out: array>; diff --git a/backends/webgpu/scripts/gen_wgsl_headers.py b/backends/webgpu/scripts/gen_wgsl_headers.py index 90293fc6cfe..a51599de7be 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -7,21 +7,39 @@ """Generate runtime/ops//_wgsl.h from each .wgsl. -Each header embeds the shader verbatim as `inline constexpr const char* +Each header embeds the shader text unchanged as `inline constexpr const char* kWGSL` plus `kWorkgroupSize` (parsed from @workgroup_size). Usage: gen_wgsl_headers.py # (re)write all _wgsl.h gen_wgsl_headers.py --check # exit 1 if any committed header is stale -Stdlib only (the devserver has no third-party pip). +A shader is treated as a template iff a sibling .yaml spec exists; the +$-block engine (preprocess/escape/generate_variant_combinations) expands one +template + a DTYPE/VEC variant matrix into the concrete per-variant headers. + +Spec parsing uses PyYAML (a declared ExecuTorch codegen dependency, mirroring +backends/vulkan/runtime/gen_vulkan_spv.py); run under the ExecuTorch dev env. """ import argparse +import copy import hashlib +import io import re import sys +from itertools import product from pathlib import Path +from typing import Any, Dict, List, Optional, Set + +import yaml +from yaml.constructor import ConstructorError +from yaml.nodes import MappingNode + +try: + from yaml import CLoader as Loader +except ImportError: + from yaml import Loader # type: ignore[assignment, misc] BACKEND_ROOT = Path(__file__).resolve().parents[1] @@ -37,6 +55,328 @@ */""" +######################################################################## +# WGSL template engine +# +# A $-block transpiler (extract_leading_whitespace / escape / preprocess) +# plus a DTYPE/VEC variant matrix (generate_variant_combinations / +# parse_template_spec) expand one template + its YAML sidecar into the +# per-variant WGSL headers. +######################################################################## + + +# WGSL type-helpers injected into preprocess's exec globals so ${...} template +# expressions can spell WGSL types (f32/f16, vec4); @group/@binding layout is +# written directly in the templates. Names mirror gen_vulkan_spv.py's +# buffer_scalar_type / buffer_gvec_type / accum_scalar_type. +def buffer_scalar_type(dtype: str) -> str: + if dtype == "half": + return "f16" + elif dtype == "float": + return "f32" + return dtype + + +def buffer_gvec_type(dtype: str, n: int) -> str: + if n == 1: + return buffer_scalar_type(dtype) + return f"vec{n}<{buffer_scalar_type(dtype)}>" + + +def accum_scalar_type(dtype: str) -> str: + # The float family (incl. half) accumulates in f32 -- f16 accumulation is + # numerically unsafe on target GPUs. Mirrors gen_vulkan_spv.py's + # accum_scalar_type (half -> rgba16f -> "float" there is the same intent). + if dtype in ("half", "float"): + return "f32" + return buffer_scalar_type(dtype) + + +WGSL_HELPERS: Dict[str, Any] = { + "buffer_scalar_type": buffer_scalar_type, + "buffer_gvec_type": buffer_gvec_type, + "accum_scalar_type": accum_scalar_type, +} + + +# https://github.com/google/XNNPACK/blob/master/tools/xngen.py +def extract_leading_whitespace(line: str) -> str: + match = re.match(r"\s*", line) + return match.group(0) if match else "" + + +# https://github.com/google/XNNPACK/blob/master/tools/xngen.py +def escape(line: str) -> str: + output_parts = [] + while "${" in line: + start_pos = line.index("${") + end_pos = line.index("}", start_pos + 2) + if start_pos != 0: + output_parts.append('"' + line[:start_pos].replace('"', '\\"') + '"') + output_parts.append("str(" + line[start_pos + 2 : end_pos] + ")") + line = line[end_pos + 1 :] + if line: + output_parts.append('"' + line.replace('"', '\\"') + '"') + return " + ".join(output_parts) + + +# https://github.com/google/XNNPACK/blob/master/tools/xngen.py +def preprocess( + input_text: str, variables: Dict[str, Any], input_path: str = "codegen" +) -> str: + # Workaround to handle source files using \ to extend mecros to a new line + input_text = re.sub(r"\\$", r"\\\\", input_text, flags=re.MULTILINE) + + input_lines = input_text.splitlines() + python_lines = [] + + blank_lines = 0 + + last_indent = "" + + # List of tuples (total_index, python_indent) + indent_stack = [("", "")] + + # Indicates whether this is the first line inside Python + # code block (i.e. for, while, if, elif, else) + python_block_start = True + for input_line in input_lines: + if input_line == "": + blank_lines += 1 + continue + # Skip lint markers. + if "LINT" in input_line: + continue + + input_indent = extract_leading_whitespace(input_line) + if python_block_start: + assert input_indent.startswith(last_indent) + extra_python_indent = input_indent[len(last_indent) :] + python_indent = indent_stack[-1][1] + extra_python_indent + indent_stack.append((input_indent, python_indent)) + assert input_indent.startswith(indent_stack[-1][0]) + else: + while not input_indent.startswith(indent_stack[-1][0]): + del indent_stack[-1] + python_block_start = False + + python_indent = indent_stack[-1][1] + stripped_input_line = input_line.strip() + if stripped_input_line.startswith("$") and not stripped_input_line.startswith( + "${" + ): + if stripped_input_line.endswith(":"): + python_block_start = True + while blank_lines != 0: + python_lines.append(python_indent + "print(file=OUT_STREAM)") + blank_lines -= 1 + python_lines.append(python_indent + stripped_input_line.replace("$", "")) + else: + assert input_line.startswith(python_indent) + while blank_lines != 0: + python_lines.append(python_indent + "print(file=OUT_STREAM)") + blank_lines -= 1 + python_lines.append( + python_indent + + "print(%s, file=OUT_STREAM)" + % escape(input_line[len(python_indent) :]) + ) + last_indent = input_indent + + while blank_lines != 0: + python_lines.append(python_indent + "print(file=OUT_STREAM)") + blank_lines -= 1 + + exec_globals = dict(variables) + output_stream = io.StringIO() + exec_globals["OUT_STREAM"] = output_stream + + python_bytecode = compile("\n".join(python_lines), input_path, "exec") + exec(python_bytecode, exec_globals) + + return output_stream.getvalue() + + +# https://gist.github.com/pypt/94d747fe5180851196eb +class UniqueKeyLoader(Loader): + def construct_mapping(self, node, deep=False): # type: ignore[no-untyped-def] + if not isinstance(node, MappingNode): + raise ConstructorError( + None, + None, + f"expected a mapping node, but found {node.id}", + node.start_mark, + ) + mapping = {} + for key_node, value_node in node.value: + key = self.construct_object(key_node, deep=deep) # type: ignore[no-untyped-call] + try: + hash(key) + except TypeError as e: + raise ConstructorError( + "while constructing a mapping", + node.start_mark, + "found unacceptable key ", + key_node.start_mark, + ) from e + # check for duplicate keys + if key in mapping: + raise ConstructorError( + "while constructing a mapping", + node.start_mark, + "found duplicate key", + key_node.start_mark, + ) + value = self.construct_object(value_node, deep=deep) # type: ignore[no-untyped-call] + mapping[key] = value + return mapping + + +def generate_variant_combinations( # noqa: C901 + iterated_params: Dict[str, Any], + exclude_params: Optional[Set[str]] = None, +) -> List[Any]: + if exclude_params is None: + exclude_params = set() + all_iterated_params = [] + for param_name, value_list in iterated_params.items(): + if re.match(r"^combination\d*$", param_name): + param_values = [] + param_names = value_list["parameter_names"] + combos = value_list["combos"] + for combo in combos: + parameter_values = combo["parameter_values"] + if "suffix" in combo: + suffix = combo["suffix"] + else: + suffix = "" + for param_value in parameter_values: + if len(str(param_value)) > 0: + suffix += "_" + str(param_value) + suffix = suffix[1:] + param_values.append((param_names, suffix, parameter_values)) + + all_iterated_params.append(param_values) + + elif param_name not in exclude_params: + param_values = [] + for value in value_list: + if "RANGE" in value: + value_range = value["RANGE"] + suffix = value.get("SUFFIX", "") + if isinstance(value_range, list) and len(value_range) == 2: + for i in range(value_range[0], value_range[1] + 1): + curr_suffix = suffix + "_" + str(i) if suffix else str(i) + param_values.append((param_name, curr_suffix, i)) + else: + raise ValueError( + f"{value['RANGE']} is not a valid range. Must be in format [start, end] (inclusive)." + ) + + elif "VALUE" in value: + suffix = value.get("SUFFIX", value["VALUE"]) + if value["VALUE"] in ["int", "uint"]: + raise ValueError( + f"Use int32 or uint32 instead of {value['VALUE']}" + ) + param_values.append((param_name, suffix, value["VALUE"])) + + else: + raise KeyError( + "Parameter must be 'VALUE: string' or 'RANGE: [a, b]'" + ) + + all_iterated_params.append(param_values) + + return list(product(*all_iterated_params)) + + +def parse_template_spec(yaml_path) -> Dict[str, List[Dict[str, Any]]]: # noqa: C901 + """Parse a .yaml variant spec into {template_name: [expanded + per-variant param dicts]}. PyYAML with a dup-key-rejecting UniqueKeyLoader + (mirrors gen_vulkan_spv.py).""" + shader_template_params: Dict[str, List[Dict[str, Any]]] = {} + with open(yaml_path) as f: + contents = yaml.load(f, Loader=UniqueKeyLoader) + for template_name, params_dict in contents.items(): + if template_name in shader_template_params: + raise KeyError(f"{template_name} params file is defined twice") + + default_params = params_dict["parameter_names_with_default_values"] + params_names = set(default_params.keys()).union({"NAME"}) + + shader_template_params[template_name] = [] + + default_iterated_params = params_dict.get("generate_variant_forall", None) + + reserved_keys = { + "generate_variant_forall", + } + + for variant in params_dict["shader_variants"]: + default_iterated_params_names = set( + default_iterated_params.keys() + if default_iterated_params is not None + else {} + ) + variant_params_names = set(variant.keys()) + + invalid_keys = ( + variant_params_names + - default_iterated_params_names + - params_names + - reserved_keys + ) + if invalid_keys: + raise ValueError(f"unknown variant key(s): {sorted(invalid_keys)}") + + iterated_params = variant.get( + "generate_variant_forall", default_iterated_params + ) + + if iterated_params is not None: + variant_combinations = generate_variant_combinations( + iterated_params, variant_params_names + ) + + for combination in variant_combinations: + default_params_copy = copy.deepcopy(default_params) + for key in variant: + if key not in reserved_keys: + default_params_copy[key] = variant[key] + + variant_name = variant["NAME"] + + for setting in combination: + param_names = setting[0] + suffix = setting[1] + param_values = setting[2] + if isinstance(param_names, list): + for param_name, param_value in zip( + param_names, param_values + ): + default_params_copy[param_name] = param_value + else: + default_params_copy[param_names] = param_values + + if len(str(suffix)) > 0: + variant_name = f"{variant_name}_{suffix}" + + default_params_copy["NAME"] = variant_name + default_params_copy["VARIANT_NAME"] = variant["NAME"] + + shader_template_params[template_name].append(default_params_copy) + else: + default_params_copy = copy.deepcopy(default_params) + for key in variant: + if key not in reserved_keys: + default_params_copy[key] = variant[key] + + shader_template_params[template_name].append(default_params_copy) + + return shader_template_params + + def symbol_base(stem: str) -> str: """snake_case shader stem -> PascalCase symbol base (binary_add -> BinaryAdd).""" return "".join(part.capitalize() for part in stem.split("_")) @@ -88,12 +428,27 @@ def embedded_sha256(header_text: str) -> str: return m.group(1) if m else "" -def render_header(wgsl_path, wgsl_text: str) -> str: - """Render the full _wgsl.h text for a shader (shader embedded verbatim).""" +def render_header( + name_or_path, wgsl_text: str, provenance_stem: Optional[str] = None +) -> str: + """Render the full _wgsl.h text for a shader (shader embedded unchanged). + + Two call forms: + - render_header(wgsl_path, wgsl_text): the plain, non-templated shaders -- + the symbol base and the `// @generated from` filename both derive from + Path(wgsl_path).stem. + - render_header(name, wgsl_text, provenance_stem): `name` is an expanded + variant name that drives the emitted symbols; `provenance_stem` is the + template stem cited in the `// @generated from` line. + """ + if provenance_stem is None: + name = Path(name_or_path).stem + provenance_stem = name + else: + name = name_or_path if ')"' in wgsl_text: raise ValueError('shader contains )" which would close the R"( literal') - stem = Path(wgsl_path).stem - base = symbol_base(stem) + base = symbol_base(name) x, y, z = parse_workgroup_size(wgsl_text) head = [ @@ -105,7 +460,7 @@ def render_header(wgsl_path, wgsl_text: str) -> str: "", "namespace executorch::backends::webgpu {", "", - f"// @generated from {stem}.wgsl - DO NOT EDIT.", + f"// @generated from {provenance_stem}.wgsl - DO NOT EDIT.", f"// wgsl-sha256: {wgsl_sha256(wgsl_text)}", f'inline constexpr const char* k{base}WGSL = R"(', ] @@ -127,6 +482,36 @@ def discover(): return sorted((BACKEND_ROOT / "runtime/ops").glob("**/*.wgsl")) +def headers_for_shader(wgsl): + """Yield (header_path, rendered_text) pairs for one shader source. + + A shader is a template iff a sibling .yaml spec exists: each expanded + variant emits its own _wgsl.h (the provenance line cites the template + stem). Otherwise the shader is embedded unchanged into _wgsl.h. + """ + stem = wgsl.stem + text = wgsl.read_text() + spec_path = wgsl.with_name(stem + ".yaml") + if spec_path.exists(): + spec = parse_template_spec(spec_path) + if list(spec.keys()) != [stem]: + raise ValueError( + f"{spec_path.name}: top-level key must be '{stem}', got {list(spec.keys())}" + ) + for variant_params in spec[stem]: + name = variant_params["NAME"] + expanded = preprocess(text, {**WGSL_HELPERS, **variant_params}) + header = wgsl.with_name(name + "_wgsl.h") + yield header, render_header(name, expanded, stem) + else: + if "$if " in text or "${" in text: + raise ValueError( + f"shader uses $if/${{ templating but has no sibling {stem}.yaml spec" + ) + header = wgsl.with_name(stem + "_wgsl.h") + yield header, render_header(stem, text, stem) + + def _report_drift(missing, stale) -> None: """Print the --check report for missing/stale committed headers.""" if missing: @@ -152,20 +537,23 @@ def main(argv=None) -> int: missing = [] errors = [] for wgsl in discover(): - wgsl_text = wgsl.read_text() try: - want = render_header(wgsl, wgsl_text) - except ValueError as e: + rendered = list(headers_for_shader(wgsl)) + # A malformed spec raises yaml.YAMLError (incl. UniqueKeyLoader's + # ConstructorError) / ValueError / KeyError from parse_template_spec, and + # a malformed template raises AssertionError from preprocess; catch them + # all so a bad shader is a clean --check report, not a traceback. + except (ValueError, KeyError, AssertionError, yaml.YAMLError) as e: errors.append(f"{wgsl.relative_to(BACKEND_ROOT)}: {e}") continue - header = wgsl.with_name(wgsl.stem + "_wgsl.h") - # Full-content compare (not just the sha) catches generator-logic drift too. - if header.exists() and header.read_text() == want: - continue - if args.check: - (missing if not header.exists() else stale).append(header) - else: - header.write_text(want) + for header, want in rendered: + # Full-content compare (not just the sha) catches generator-logic drift too. + if header.exists() and header.read_text() == want: + continue + if args.check: + (missing if not header.exists() else stale).append(header) + else: + header.write_text(want) if errors: print("Cannot generate header (malformed shader):") diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 283279e4fb5..7d32a230459 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -95,11 +95,13 @@ def test_committed_headers_match_generator(self) -> None: wgsls = g.discover() self.assertGreater(len(wgsls), 0, "no .wgsl shaders discovered") for wgsl in wgsls: - want = g.render_header(wgsl, wgsl.read_text()) - got = wgsl.with_name(wgsl.stem + "_wgsl.h").read_text() - self.assertEqual( - got, want, f"{wgsl.stem}_wgsl.h stale; run scripts/gen_wgsl_headers.py" - ) + # headers_for_shader handles both verbatim shaders and templates + # (a template emits one header per expanded variant). + for header, want in g.headers_for_shader(wgsl): + got = header.read_text() + self.assertEqual( + got, want, f"{header.name} stale; run scripts/gen_wgsl_headers.py" + ) def test_parse_workgroup_allows_space(self) -> None: # @workgroup_size (64) — the spec-legal spaced form must still parse.