[ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen

[JCNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] Steel q4gsw prefill GEMM — test coverage #20731
• [ExecuTorch][WebGPU] Add 256-thread "steel" q4gsw prefill GEMM #20730
• [ExecuTorch][WebGPU] Request shader-f16 device feature (enablement, fail-open) #20729
• -> [ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen #20728
• [ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup #20727

Additive coverage for the template engine added in the previous diff.

Key additions (test/test_wgsl_codegen.py, WgslTemplateEngineTest):

• preprocess $if/$else branch selection + ${...} substitution + guarded-body indentation.
• the DTYPE x VEC variant matrix and SUFFIX rules ("" suppresses; defaults to str(VALUE)).
• parse_template_spec — minimal-spec expansion, duplicate-key rejection (via the UniqueKeyLoader), top-level-key validation.
• the 3 WGSL type-helpers (buffer_scalar_type / buffer_gvec_type / accum_scalar_type).
• a byte-identity round-trip reproducing both committed rms_norm headers from the template + rms_norm.yaml.
• a source-of-truth sync-lock asserting the shared $-block transpiler helpers and the UniqueKeyLoader stay character-identical to gen_vulkan_spv.py (guards against silent drift).
  @exported-using-ghexport

Differential Revision: D110660006

Differential Revision: D110660006

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20728

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Cancelled Job, 1 Unrelated Failure

As of commit d3a4296 with merge base 4af91c3 ():

NEW FAILURE - The following job has failed:

• pull / test-binary-size-linux-gcc / linux-job (gh)
  RuntimeError: Command docker exec -t a3c2455f7fae6c9858c97d5ce0a13cb0049cd49667fa84d8f40753e337d84e29 /exec failed with exit code 1

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / unittest-nxp-neutron / linux-job (gh)
  ##[error]The operation was canceled.

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: JCNTH / name: Julian Ng-Thow-Hing (2a0675a, 92e5a83)

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[JCNTH]

@claude throoughly review, check for overengineering, correctness, and areaas for modularization

[claude]

Claude finished @JCNTH's task in 2m 40s —— View job

---

Review: WGSL shader-variant codegen unit tests

• Read the test file
• Read the codegen source (gen_wgsl_headers.py) under test
• Cross-check the sync-lock against gen_vulkan_spv.py
• Assess correctness, overengineering, modularization

Solid, well-targeted test suite. The names are descriptive, the docstrings capture the why (the DD-009 enable f16 note and the raw-string-terminator case are especially good), and every assertion I traced lines up with the actual API in gen_wgsl_headers.py. I verified the three shared helpers are byte-identical between the two files, and all files the tests reference (rms_norm.wgsl/.json/.h) exist. Verdict: approve with minor suggestions — nothing blocking.

Note: I reviewed by inspection + cross-file diffing; the sandbox blocked me from actually running pytest, so I did not get a green run locally.

---

Correctness

• _function_source (lines 36–61) is the one piece with real fragility. Its body-end heuristic — "stop at the next non-blank column-0 line" — is correct for the current layout (escape is followed by a # comment, preprocess by class SPVGenerator:, both at column 0), but it will silently truncate if any helper ever gains a line that starts at column 0 inside the body (e.g. a multi-line string, or a ) closing a call at column 0). For a test whose whole purpose is guarding against drift, the parser guarding it shouldn't itself be position-sensitive. See the modularization note for a sturdier approach.

• The two-phase signature skip (lines 50–53) correctly handles preprocess's multi-line signature whose ) -> str: sits at column 0 — good catch by the author. Worth noting the loop assumes the first line ending in : is the signature terminator; a default value like x: Dict = {} ending a wrapped signature line would fool it. Not a problem for these three functions, but it's the same class of latent fragility.

Overengineering

• Reasonable overall — this isn't gold-plated. The single candidate is the sync-lock's hand-rolled text parser: ~26 lines of column heuristics to enforce char-identity of 3 functions. Since gen_wgsl_headers.py already imports these functions, the webgpu side can just use inspect.getsource(getattr(g, fn)) and only the (un-importable) Vulkan side needs source extraction — and that extraction is more robustly done with ast than regex/column-counting. That removes the fragility and shrinks the helper.

• Mild redundancy: test_rms_norm_template_roundtrip_byte_identical (416) overlaps test_committed_headers_match_generator (134) — the committed rms_norm headers are already validated through headers_for_shader, which runs the same template expansion. The round-trip test earns its keep by exercising preprocess + render_header directly (bypassing spec parsing) and pinning the concrete variant names, so I'd keep it — just calling out the overlap.

Modularization

• Suggested: extract via ast instead of column heuristics. ast.parse(text) + walk for FunctionDef with the target name + ast.get_source_segment(text, node) gives the exact source span with no indentation assumptions, and works identically for both files. This replaces _function_source entirely and is pure stdlib (consistent with the generator's stdlib-only constraint). Fix this →

• The two test classes (WgslCodegenTest / WgslTemplateEngineTest) and the module-level loader are cleanly separated. _write_spec is an appropriate local helper. No changes needed here.

Coverage gaps (optional, not blocking)

generate_variant_combinations/parse_template_spec are inherited from the Vulkan engine and have branches the tests don't reach. If webgpu will exercise them, worth a line each:

• The int/uint rejection (gen_wgsl_headers.py:240-243) — a one-line assertRaises(ValueError) test.
• The RANGE expansion (lines 226-236), including the _1/_2 suffix naming — currently only VALUE is tested.
• The invalid variant-key assert in parse_template_spec (line 292) — the top-level-key mismatch is covered, but not an unknown key inside a shader_variants entry.
• The combination/combinationN branch (lines 205-221) is untested. Likely out of scope for webgpu today; a log-style comment noting it's intentionally uncovered would prevent a false sense of completeness.

Nit

• test_render_header_shape:111 asserts assertNotIn("kUpdateCacheWorkgroupSize ", h) (trailing space) to reject an untyped symbol — clever, but a touch cryptic; a short comment (# no untyped WorkgroupSize symbol, only X/Y/Z) would help the next reader.

---

None of the above blocks the PR — the suggestions are hardening the drift-guard and closing a few inherited-branch gaps. Nice work on the byte-identity and DD-009 regression cases.
· branch gh/JCNTH/2/head