From ecdd8230b62d1d223205e1cedff1be760fdc145b Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Sat, 4 Jul 2026 18:19:43 -0700 Subject: [PATCH 1/4] Update [ghstack-poisoned] --- .../webgpu/runtime/ops/rms_norm/rms_norm.json | 17 + .../webgpu/runtime/ops/rms_norm/rms_norm.wgsl | 75 ++-- .../runtime/ops/rms_norm/rms_norm_vec4.wgsl | 74 ---- .../runtime/ops/rms_norm/rms_norm_vec4_wgsl.h | 2 +- backends/webgpu/scripts/gen_wgsl_headers.py | 377 +++++++++++++++++- backends/webgpu/test/test_wgsl_codegen.py | 12 +- 6 files changed, 437 insertions(+), 120 deletions(-) create mode 100644 backends/webgpu/runtime/ops/rms_norm/rms_norm.json delete mode 100644 backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4.wgsl diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm.json b/backends/webgpu/runtime/ops/rms_norm/rms_norm.json new file mode 100644 index 00000000000..651235cb325 --- /dev/null +++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm.json @@ -0,0 +1,17 @@ +{ + "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.wgsl b/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl index 4bd5618596f..9d7ce1d138d 100644 --- a/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl +++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl @@ -1,6 +1,6 @@ -@group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_in: array; -@group(0) @binding(2) var t_weight: array; +@group(0) @binding(0) var t_out: array<${wgsl_buffer_type(DTYPE, VEC)}>; +@group(0) @binding(1) var t_in: array<${wgsl_buffer_type(DTYPE, VEC)}>; +@group(0) @binding(2) var t_weight: array<${wgsl_buffer_type(DTYPE, VEC)}>; struct Params { num_rows: u32, @@ -12,7 +12,7 @@ struct Params { const WG_SIZE: u32 = 64u; -var shared_sum: array; +var shared_sum: array<${wgsl_accum_type()}, WG_SIZE>; fn reduce_shared(worker_id: u32) { workgroupBarrier(); @@ -29,6 +29,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 +44,33 @@ 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: ${wgsl_accum_type()} = 0.0; + $if VEC == 4: + 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; + } + $else: + var x: u32 = worker_id; + loop { + if (x >= params.row_width) { + break; + } + let v = t_in[base + x]; + 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 +78,24 @@ 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; + } + 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]; + 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_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..764647676ae 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -7,21 +7,30 @@ """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 .json spec exists; the +$-block engine (preprocess/escape/generate_variant_combinations) expands one +template + a DTYPE/VEC variant matrix into the concrete per-variant headers. + +Stdlib only (the devserver has no third-party pip; no yaml). """ import argparse +import copy import hashlib +import io +import json import re import sys +from itertools import product from pathlib import Path +from typing import Any, Dict, List, Optional, Set BACKEND_ROOT = Path(__file__).resolve().parents[1] @@ -37,6 +46,298 @@ */""" +######################################################################## +# 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 JSON 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. +def wgsl_scalar_type(dtype: str) -> str: + if dtype == "half": + return "f16" + elif dtype == "float": + return "f32" + return dtype + + +def wgsl_buffer_type(dtype: str, vec: int) -> str: + if vec == 1: + return wgsl_scalar_type(dtype) + return f"vec{vec}<{wgsl_scalar_type(dtype)}>" + + +def wgsl_accum_type() -> str: + # Accumulators stay f32 in every variant (f16 accumulation is numerically + # unsafe on target GPUs). + return "f32" + + +WGSL_HELPERS: Dict[str, Any] = { + "wgsl_scalar_type": wgsl_scalar_type, + "wgsl_buffer_type": wgsl_buffer_type, + "wgsl_accum_type": wgsl_accum_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() + + +# json object_pairs_hook that rejects duplicate keys in a spec object. +def _reject_duplicate_keys(pairs: List[Any]) -> Dict[str, Any]: + mapping: Dict[str, Any] = {} + for key, value in pairs: + if key in mapping: + raise ValueError(f"found duplicate key: {key!r}") + 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(json_path) -> Dict[str, List[Dict[str, Any]]]: # noqa: C901 + """Parse a .json variant spec into {template_name: [expanded + per-variant param dicts]}. Stdlib JSON with a dup-key-rejecting + object_pairs_hook.""" + shader_template_params: Dict[str, List[Dict[str, Any]]] = {} + with open(json_path) as f: + contents = json.load(f, object_pairs_hook=_reject_duplicate_keys) + 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 + ) + assert len(invalid_keys) == 0 + + 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 +389,25 @@ 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: 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 +419,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 +441,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 .json 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 + ".json") + 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}.json 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 +496,19 @@ def main(argv=None) -> int: missing = [] errors = [] for wgsl in discover(): - wgsl_text = wgsl.read_text() try: - want = render_header(wgsl, wgsl_text) + rendered = list(headers_for_shader(wgsl)) except ValueError 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. From 5b36ffd03a87dd7edf1c43c41f715d3754d4947b Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Sat, 4 Jul 2026 19:27:30 -0700 Subject: [PATCH 2/4] Update [ghstack-poisoned] --- backends/webgpu/scripts/gen_wgsl_headers.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/backends/webgpu/scripts/gen_wgsl_headers.py b/backends/webgpu/scripts/gen_wgsl_headers.py index 764647676ae..ccd1fc7ea25 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -289,7 +289,8 @@ def parse_template_spec(json_path) -> Dict[str, List[Dict[str, Any]]]: # noqa: - params_names - reserved_keys ) - assert len(invalid_keys) == 0 + if invalid_keys: + raise ValueError(f"unknown variant key(s): {sorted(invalid_keys)}") iterated_params = variant.get( "generate_variant_forall", default_iterated_params @@ -389,7 +390,9 @@ def embedded_sha256(header_text: str) -> str: return m.group(1) if m else "" -def render_header(name_or_path, wgsl_text: str, provenance_stem: str = None) -> str: +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: @@ -498,7 +501,10 @@ def main(argv=None) -> int: for wgsl in discover(): try: rendered = list(headers_for_shader(wgsl)) - except ValueError as e: + # A malformed spec raises ValueError/KeyError from parse_template_spec, + # and a malformed template raises AssertionError from preprocess; catch + # all three so a bad shader is a clean --check report, not a traceback. + except (ValueError, KeyError, AssertionError) as e: errors.append(f"{wgsl.relative_to(BACKEND_ROOT)}: {e}") continue for header, want in rendered: From b7baf3aa71c211a13e20785d85785aeb89f5a3d1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Sat, 4 Jul 2026 21:31:42 -0700 Subject: [PATCH 3/4] Update [ghstack-poisoned] --- .../webgpu/runtime/ops/rms_norm/rms_norm.wgsl | 32 +++++++++++++------ backends/webgpu/scripts/gen_wgsl_headers.py | 30 +++++++++-------- 2 files changed, 40 insertions(+), 22 deletions(-) diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl b/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl index 9d7ce1d138d..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<${wgsl_buffer_type(DTYPE, VEC)}>; -@group(0) @binding(1) var t_in: array<${wgsl_buffer_type(DTYPE, VEC)}>; -@group(0) @binding(2) var t_weight: array<${wgsl_buffer_type(DTYPE, VEC)}>; +$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<${wgsl_accum_type()}, WG_SIZE>; +var shared_sum: array<${accum_scalar_type(DTYPE)}, WG_SIZE>; fn reduce_shared(worker_id: u32) { workgroupBarrier(); @@ -50,7 +52,7 @@ fn main( $else: let base = row_idx * params.row_width; - var local_sq_sum: ${wgsl_accum_type()} = 0.0; + var local_sq_sum: ${accum_scalar_type(DTYPE)} = 0.0; $if VEC == 4: var x4: u32 = worker_id; loop { @@ -58,7 +60,10 @@ fn main( break; } let v = t_in[base4 + x4]; - local_sq_sum = local_sq_sum + dot(v, v); + $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: @@ -68,7 +73,10 @@ fn main( break; } let v = t_in[base + x]; - local_sq_sum = local_sq_sum + v * v; + $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; } @@ -84,7 +92,10 @@ fn main( if (x4 >= rw4) { break; } - t_out[base4 + x4] = t_in[base4 + x4] * rstd * t_weight[x4]; + $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: @@ -95,7 +106,10 @@ fn main( } let v = t_in[base + x]; let w = t_weight[x]; - t_out[base + x] = v * rstd * w; + $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; } } diff --git a/backends/webgpu/scripts/gen_wgsl_headers.py b/backends/webgpu/scripts/gen_wgsl_headers.py index ccd1fc7ea25..67729d6a8f7 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -58,8 +58,9 @@ # 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. -def wgsl_scalar_type(dtype: str) -> str: +# 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": @@ -67,22 +68,25 @@ def wgsl_scalar_type(dtype: str) -> str: return dtype -def wgsl_buffer_type(dtype: str, vec: int) -> str: - if vec == 1: - return wgsl_scalar_type(dtype) - return f"vec{vec}<{wgsl_scalar_type(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 wgsl_accum_type() -> str: - # Accumulators stay f32 in every variant (f16 accumulation is numerically - # unsafe on target GPUs). - return "f32" +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] = { - "wgsl_scalar_type": wgsl_scalar_type, - "wgsl_buffer_type": wgsl_buffer_type, - "wgsl_accum_type": wgsl_accum_type, + "buffer_scalar_type": buffer_scalar_type, + "buffer_gvec_type": buffer_gvec_type, + "accum_scalar_type": accum_scalar_type, } From b06c0d3104d0e8d1749522dcbc9e7cdce3e595a7 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Sat, 4 Jul 2026 23:18:21 -0700 Subject: [PATCH 4/4] Update [ghstack-poisoned] --- .../webgpu/runtime/ops/rms_norm/rms_norm.json | 17 ---- .../webgpu/runtime/ops/rms_norm/rms_norm.yaml | 12 +++ backends/webgpu/scripts/gen_wgsl_headers.py | 85 +++++++++++++------ 3 files changed, 72 insertions(+), 42 deletions(-) delete mode 100644 backends/webgpu/runtime/ops/rms_norm/rms_norm.json create mode 100644 backends/webgpu/runtime/ops/rms_norm/rms_norm.yaml diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm.json b/backends/webgpu/runtime/ops/rms_norm/rms_norm.json deleted file mode 100644 index 651235cb325..00000000000 --- a/backends/webgpu/runtime/ops/rms_norm/rms_norm.json +++ /dev/null @@ -1,17 +0,0 @@ -{ - "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.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/scripts/gen_wgsl_headers.py b/backends/webgpu/scripts/gen_wgsl_headers.py index 67729d6a8f7..a51599de7be 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -14,24 +14,33 @@ gen_wgsl_headers.py # (re)write all _wgsl.h gen_wgsl_headers.py --check # exit 1 if any committed header is stale -A shader is treated as a template iff a sibling .json spec exists; the +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. -Stdlib only (the devserver has no third-party pip; no yaml). +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 json 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] _SHA_RE = re.compile(r"// wgsl-sha256: ([0-9a-f]{64})") @@ -51,7 +60,7 @@ # # A $-block transpiler (extract_leading_whitespace / escape / preprocess) # plus a DTYPE/VEC variant matrix (generate_variant_combinations / -# parse_template_spec) expand one template + its JSON sidecar into the +# parse_template_spec) expand one template + its YAML sidecar into the # per-variant WGSL headers. ######################################################################## @@ -188,14 +197,39 @@ def preprocess( return output_stream.getvalue() -# json object_pairs_hook that rejects duplicate keys in a spec object. -def _reject_duplicate_keys(pairs: List[Any]) -> Dict[str, Any]: - mapping: Dict[str, Any] = {} - for key, value in pairs: - if key in mapping: - raise ValueError(f"found duplicate key: {key!r}") - mapping[key] = value - return mapping +# 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 @@ -257,13 +291,13 @@ def generate_variant_combinations( # noqa: C901 return list(product(*all_iterated_params)) -def parse_template_spec(json_path) -> Dict[str, List[Dict[str, Any]]]: # noqa: C901 - """Parse a .json variant spec into {template_name: [expanded - per-variant param dicts]}. Stdlib JSON with a dup-key-rejecting - object_pairs_hook.""" +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(json_path) as f: - contents = json.load(f, object_pairs_hook=_reject_duplicate_keys) + 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") @@ -451,13 +485,13 @@ def discover(): def headers_for_shader(wgsl): """Yield (header_path, rendered_text) pairs for one shader source. - A shader is a template iff a sibling .json spec exists: each expanded + 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 + ".json") + spec_path = wgsl.with_name(stem + ".yaml") if spec_path.exists(): spec = parse_template_spec(spec_path) if list(spec.keys()) != [stem]: @@ -472,7 +506,7 @@ def headers_for_shader(wgsl): else: if "$if " in text or "${" in text: raise ValueError( - f"shader uses $if/${{ templating but has no sibling {stem}.json spec" + 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) @@ -505,10 +539,11 @@ def main(argv=None) -> int: for wgsl in discover(): try: rendered = list(headers_for_shader(wgsl)) - # A malformed spec raises ValueError/KeyError from parse_template_spec, - # and a malformed template raises AssertionError from preprocess; catch - # all three so a bad shader is a clean --check report, not a traceback. - except (ValueError, KeyError, AssertionError) as e: + # 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 for header, want in rendered: