[ExecuTorch][WebGPU] Generate extrema and unary shader variants

[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

Generate extrema and unary shader families

The extrema reductions and ten no-parameter unary kernels duplicated shader skeletons that could drift independently. This consolidates amax/amin behind one extrema template and abs/cos/exp/hardswish/neg/round/rsqrt/sin/sqrt/tanh behind one unary template while preserving the generated runtime payloads.

Key changes:

• Generate amax/amin from one extrema manifest.
• Generate ten unary payloads from one operator-expression manifest.
• Lock expanded bytes, registry entries, delegation, and boundary numerics.

The attempted Unary lifecycle migration is intentionally not part of the stack: its performance campaign did not produce an authoritative passing result, so the Unary builder, interface, and activation/sigmoid call sites are restored to their pre-migration bytes.

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

Differential Revision: D113979760

Differential Revision: D113979760

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 63b0192 with merge base f4a21b8 ():

NEW FAILURES - The following jobs have failed:

• pull / test-llama-runner-linux (fp32, xnnpack+quantize_kv, linux.arm64.2xlarge, executorch-ubuntu-22.04-... / linux-job (gh)
  RuntimeError: Command docker exec -t 94d1d338b02203d7cea9316c162d0b1a99f008e1742614f74b51995527e03bc7 /exec failed with exit code 1
• pull / test-models-linux (llama3_2_vision_encoder, portable, linux.4xlarge.memory) / linux-job (gh)
  RuntimeError: Command docker exec -t 46b49c34f1665cfb9f5694a6f146b1ea17448247ccbabf4b56693f170d3df036 /exec failed with exit code 1

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 4m 50s —— View job

---

Review: Generate extrema shader variants

• Read the full diff vs base
• Inspect shader template + YAML codegen
• Check handler/registry relocations
• Review test coverage
• Post findings

The core change is clean and correct. extrema.wgsl + extrema.yaml follow the existing template convention exactly (mirrors binary_op, to_copy, etc.), the embedded wgsl-sha256 values are unchanged (35fc…/8cb6…), confirming the shader text is byte-identical to the old amax.wgsl/amin.wgsl. The handler edits are include-path-only, the header relocation to runtime/ops/extrema/ matches the binary_op pattern (handler in amax/, generated header in extrema/), and the CMake runtime/ops/*/*.cpp glob means no build-file edit is needed. The sign_trap/tie fixtures use exactly-representable fp32 values (/8.0, /16.0), so the atol=0/rtol=0 comparisons are safe, and the dim=1 "posdim" cases correctly still resolve to last-dim on 2D inputs. Good coverage of the boundary widths (63/64/65, 255/256/257) around wg_size=256.

A few things worth reconsidering, in priority order:

1. test_config_contract is a tautological change-detector — recommend removing

test_reduce.py:217 asserts EXTREMA_CONFIGS == (<the same tuple retyped verbatim>). It can only ever fail when someone edits the constant, and then it fails by design with no bug caught — it just forces the same edit in two places. This is pure maintenance cost. If the intent is to guard the contract rather than the literal, assert the invariants that actually matter instead:

names = [c[0] for c in EXTREMA_CONFIGS]
self.assertEqual(len(names), len(set(names)))                     # unique names
self.assertTrue(all(c[4] in {"default", "sign_trap", "tie"} for c in EXTREMA_CONFIGS))
self.assertTrue({63, 64, 65, 255, 256, 257} <= {c[1][-1] for c in EXTREMA_CONFIGS})  # boundary widths

Fix this →

2. Byte-hashing the hand-written Reduce.cpp handlers is brittle

test_wgsl_codegen.py:1014 pins sha256(Reduce.cpp) for both handlers. Those files are not generated — any whitespace, comment, or unrelated refactor to a handler breaks a codegen test, with a hash diff that tells the next person nothing about what changed or why. The generated-artifact hashes (the WGSL bodies and registry_path()) are legitimate codegen contracts; hashing handwritten source isn't. If the goal is "the handler wires up the relocated header," assert that directly and legibly:

self.assertIn(f"ops/extrema/{name}_wgsl.h", handler.read_text())

Fix this →

3. Modularization: the two Reduce.cpp handlers are still ~155-line near-verbatim duplicates

amax/Reduce.cpp and amin/Reduce.cpp differ only in identifiers (amax↔amin, param/symbol names) — the dispatch, resize hook, and param layout are identical. This PR unified the shaders but the handlers remain fully copied. This is consistent with the existing div/sub split (also duplicated handlers under the CMake glob), so it's not a blocker — but if the motivation for the PR is "these could drift independently," the same argument applies to the handlers, and a shared reduce_per_row_impl(graph, args, kXxxWGSL, "amax") helper would eliminate the larger duplication. Reasonable to defer, but worth a note in the PR description that the handler duplication is intentional/out-of-scope.

4. Minor: duplicated (op, input_class) → generator dispatch

cases.py:501 builds a generators dict that mirrors the same dispatch already encoded in _extrema_input (test_reduce.py:120). Two sources of truth for the same mapping. Since InputSpec.gen wants a shape → tensor callable, functools.partial(_extrema_input, op, input_class) would let both call sites share the one dispatch.

5. Nit: stale filename in the static_assert comment

amax/Reduce.cpp:24 (and the amin mirror) still reads "amax.wgsl partials[] is sized 256…", but the source is now extrema.wgsl. Cheap to update to extrema.wgsl while you're touching these files.

---

None of these are correctness bugs — the runtime math, dispatch, and payloads are genuinely preserved. #1 and #2 are the ones I'd actually change before merge; the rest are optional.

Note: I reviewed statically and could not execute gen_wgsl_headers.py --check or the test suite in this environment (would need the ExecuTorch dev env / added --allowedTools), so the byte-identical claims are verified by reading the embedded hashes rather than by running codegen.
· branch gh/JCNTH/200/head