[ExecuTorch][WebGPU] Make WGSL generation rollback-safe

[JCNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] Preserve Linear and Q4 embedding params across dynamic resize #21483
• [ExecuTorch][WebGPU] Generate logical and arithmetic binary shader variants #21451
• [ExecuTorch][WebGPU] Generate extrema and unary shader variants #21450
• [ExecuTorch][WebGPU] Generate typed to_copy shader variants #21449
• -> [ExecuTorch][WebGPU] Make WGSL generation rollback-safe #21448
• [ExecuTorch][WebGPU] Add Qwen3 K16 attention and fp16 KV cache #21136
• [ExecuTorch][WebGPU] Add HuggingFace rotate-half RoPE #21135
• [ExecuTorch][WebGPU] Repair dynamic SDPA routing, add attestation #21134
• [ExecuTorch][WebGPU] Select emdawnwebgpu for Emscripten builds #21133
• [ExecuTorch][WebGPU] Add Llama K16 online causal attention #21132
• [ExecuTorch][WebGPU] Replace broad QKV fusion with BK64 kernel #21131
• [ExecuTorch][WebGPU] Route exact prefill shapes to the BK64 GEMM #21130
• [ExecuTorch][WebGPU] Fuse the SwiGLU elementwise chain #21129
• [ExecuTorch][WebGPU] Coordinate dynamic dispatch routes #21128
• [ExecuTorch][WebGPU] Add certified scoped output suppression #21127

Make WGSL publication rollback-safe without changing generated runtime output

Generated headers and the registry were previously written incrementally, so later validation, interruption, or write failure could leave a partial generated tree. This change stages and validates the complete output set before sequential publication, rejects path/name/C++ symbol collisions, and restores replaced outputs on failure or interruption.

The family model follows Vulkan's YAML/template generation. The WebGPU publisher adds staging, backup, rollback, and temporary-file cleanup; it does not claim filesystem-atomic multi-file replacement.

Key changes:

• gen_wgsl_headers.py — stage, validate, publish, restore on failure, and clean temporary outputs.
• test_wgsl_codegen.py — cover collisions, orphan repair, interruption, rollback, and idempotence.

All generated shader bytes and runtime registry entries remain unchanged.

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D113979658

Differential Revision: D113979658

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 8 Pending

As of commit b3701c7 with merge base f4a21b8 ():
💚 Looks good so far! There are no failures yet. 💚

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

[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 Double check for need of changes, modulariziation opportunity, etc

[claude]

Claude finished @JCNTH's task in 3m 27s —— View job

---

Review: transactional WGSL generation

• Read gen_wgsl_headers.py changes
• Read test_wgsl_codegen.py
• Analyze for bugs, correctness, modularization
• Post review

Verdict: The transactional logic is correct and unusually well-tested — I found no correctness bugs. Below are observations on design/modularization and a few minor polish items. (Note: I couldn't execute pytest in this sandbox — permission denied — so this is a static review.)

---

Correctness — looks sound

I traced the stage → publish → rollback paths against the tests and they hold up:

• _stage_outputs stages temps before touching any destination; on OSError it cleans temps and returns errors (no destinations replaced yet, so no rollback needed); on other BaseException (interrupt) it cleans temps and re-raises. Correct — nothing to roll back at stage time.
• _publish_outputs conservatively treats an interrupt during os.replace as possibly-replaced (appends to replaced) so rollback covers it; restore is idempotent for the not-actually-replaced case. Good.
• _rollback_outputs restores in reversed(replaced) order, keeps going past individual restore failures, and reports them (test_multiple_rollback_errors_do_not_stop_later_restores). Solid.
• Mode preservation for existing files / 0o644 for new files is handled and tested.

Behavior change worth calling out

Orphan headers now block normal generation, not just --check. In main, if orphans: return 1 runs before _publish_outputs, so a stray *_wgsl.h with no source prevents regenerating any other (legitimately stale) header until the user removes it by hand. Previously orphans weren't detected at all. This is a defensible "never delete, surface it" stance and is tested (test_orphans_are_sorted_reported_and_never_deleted) — just confirm it's the intended UX, since it makes the generator stricter than before. A one-line note in the docstring/--help ("generation refuses to run while orphan headers exist") would save a confused user.

Modularization opportunity (the main one)

The entire transaction layer — _OriginalOutput, _stage_bytes, _cleanup_temporaries, _stage_outputs, _rollback_outputs, _publish_outputs — is fully generic: it's an atomic multi-file publish over Dict[Path, bytes] with zero WGSL knowledge. It's ~110 lines that dominate the diff and are conceptually separate from "render WGSL headers."

Options, in order of how much I'd lean toward them:

1. Leave it inline (simplest, fewest files) — fine given it's used in exactly one place today. CLAUDE.md's "prefer simpler/concise" arguably favors this.
2. Extract to a small sibling module (e.g. scripts/_atomic_publish.py) — separates the generic transaction concern from codegen, makes it unit-testable in isolation, and shrinks gen_wgsl_headers.py back to its domain logic. Worth it only if you expect to reuse it (e.g. the sibling Vulkan generator).

I'd only do (2) if reuse is likely; otherwise the inline form is fine. Extract this →

Minor polish (all optional)

1. Undocumented except OSError: raise split in _publish_outputs — the reason it's separated from except BaseException (OSError = replace failed, so don't mark it replaced; interrupt = might have succeeded, so do) is subtle and non-obvious. One comment would save the next reader a few minutes. Fix this →

2. registry_entries() now renders every header (via _collect_header_outputs) just to return entries, whereas the old version discarded the text. The main path is fine (single render through collect_outputs), but the public registry_entries() helper and test_registry_entries_match_concrete_headers now each trigger a full render. Not a bug — just note that registry_entries() is heavier than its name implies; a caller that only wants names pays for full expansion.

3. sorted(changed) is computed twice (once in _stage_outputs, once in _publish_outputs). Consistent and deterministic so pairing is safe, but you could sort once in _publish_outputs and pass the ordered list down.

4. Redundant name-collision check: _collect_header_outputs rejects duplicate registry names/symbols, and render_registry independently re-checks len(names) != len(set(names)). Both are defensive and tested — fine to keep, just noting the overlap.

Test coverage

Excellent — collisions (path/name/symbol), staging interrupt, replace failure with restore, multi-error rollback, orphan reporting, idempotent second run with no IO, and interrupt-then-repair are all covered. The test_generation_renders_once_and_second_run_does_no_io test is a particularly nice guard on both the single-render and no-write-when-clean invariants.

Nothing here is blocking. My one substantive ask is confirming the orphan-blocks-generation behavior is intended (and documenting it); the rest is optional cleanup.
· branch gh/JCNTH/198/head