export/_whisper.py

"""Whisper base building blocks for recipes/whisper-base.py.

Everything here is built from the pinned openai/whisper-base checkpoint loaded through transformers
(WhisperForConditionalGeneration); the program wrappers reuse its nn.Modules (weights are shared, not
copied) and only re-express the forward pass so that it lowers to one WebGPU delegate per program.

Program contract (same split and I/O as the playground's previous whisper-base export):

  mel            forward(frames[1,3000,400] f32) -> power_mel[1,3000,80] f32
                 frames: 30 s of 16 kHz audio, zero-padded to 480000 samples, reflect-padded by 200 on
                 both ends, 3000 windows of 400 samples at hop 160 (= whisper.js buildFrames).
                 power_mel = |rfft(hann_periodic * frame)|^2 @ mel_filters (no log; see host_log_mel).
  encoder        forward(log_mel[1,80,3000] f32) -> encoder_hidden[1,1500,512] f32   (fp32 weights)
  cross_kv       forward(encoder_hidden[1,1500,512]) -> 12 x [1,8,1500,64] f32, order k0,v0,...,k5,v5
                 (per decoder layer: encoder_attn.k_proj / v_proj, heads-first [B,H,S,D])
  decoder_split  forward(tok_embeds[1,64,512], k0, v0, ..., k5, v5) -> hidden[1,64,512]
                 tok_embeds: token embeddings WITHOUT positions for positions 0..64 (zero rows beyond
                 the live prefix); the program adds positions 0..63, runs causal self-attention over
                 the 64 rows and returns final-LayerNorm hidden states (row p predicts token p+1).
  decode_stream  forward(embed[1,1,512], input_pos[1] i64, k0, v0, ..., k5, v5) -> hidden[1,1,512]
                 embed = token embedding + positional embedding of `input_pos` (added on the host);
                 self-attention K/V cache [1,448,8,64] x 6 layers x {k,v} lives inside the program
                 (mutable buffers owned by the delegate); feed positions 0,1,2,... in order.
  lmhead         forward(hidden[1,1,512]) -> logits[1,1,51872]  (tied to decoder.embed_tokens; rows
                 51865..51871 are padding with logit -1e9, see LMHeadProgram)

Precision (whisper-base.py --precision):
  q4   (previous-export parity) cross_kv / decoder_split / decode_stream / lmhead linears are 4-bit groupwise
       weight-only (torchao IntxWeightOnlyConfig int4, group 64, symmetric, "affine" qparams ->
       et_vk.linear_q4gsw); int8 token-embedding table. mel and encoder are fp32.
  fp32 every program fp32 (aten.linear on 2-D views); fp32 token-embedding table.

Host-side files (reproduced from the checkpoint by the recipe, consumed by the playground JS):
  whisper_embed.i8.bin         int8 [51865,512]  round(embed_tokens / scale), scale = rowwise absmax/127  (q4)
  whisper_embed.scale.f32.bin  f32  [51865]      the per-row scales                                       (q4)
  whisper_embed.f32.bin        f32  [51865,512]  embed_tokens                                             (fp32)
  whisper_pos_embed.f32.bin    f32  [448,512]    decoder.embed_positions.weight
  vocab.json                   upstream tokenizer vocab (byte-level BPE, detokenization only)
"""

import copy
import io

import numpy as np
import torch
import torch.nn.functional as F

SAMPLE_RATE = 16000
N_FFT = 400
HOP = 160
FRAMES = 3000
N_SAMPLES = FRAMES * HOP
PAD = N_FFT // 2
N_FREQ = N_FFT // 2 + 1
N_MELS = 80
ENC_LEN = 1500
D_MODEL = 512
N_HEADS = 8
HEAD_DIM = 64
N_LAYERS = 6
MAX_DEC = 64  # decoder_split rows; the playground's decode budget (prefix + generated)
KV_LEN = 448  # decode_stream self-attention cache length (= max_target_positions)
VOCAB = 51865
VOCAB_PADDED = 51872  # lmhead rows (multiple of 8, = the previous export)
PAD_LOGIT = -1e9

TOK = {
    "sot": 50258,
    "eot": 50257,
    "en": 50259,
    "transcribe": 50359,
    "notimestamps": 50363,
}
PREFIX_EN = [TOK["sot"], TOK["en"], TOK["transcribe"], TOK["notimestamps"]]


# --------------------------------------------------------------------------------------------
# Host-side math that the playground does in JS (kept here so validation runs the same steps).
# --------------------------------------------------------------------------------------------


def build_frames(pcm):
    """Mirror of whisper.js buildFrames(): [3000, 400] float32 windows (no window function)."""
    pcm = np.asarray(pcm, dtype=np.float32)
    base = np.zeros(N_SAMPLES + 2 * PAD, dtype=np.float32)
    n = min(pcm.size, N_SAMPLES)
    base[PAD : PAD + n] = pcm[:n]
    for i in range(PAD):
        base[PAD - 1 - i] = base[PAD + 1 + i]
        base[PAD + N_SAMPLES + i] = base[PAD + N_SAMPLES - 2 - i]
    idx = np.arange(FRAMES)[:, None] * HOP + np.arange(N_FFT)[None, :]
    return base[idx]


def host_log_mel(power_mel):
    """Host step between the mel and encoder programs (WhisperFeatureExtractor normalization).

    power_mel [1,3000,80] (or [3000,80]) -> log_mel [1,80,3000]:
      x = log10(max(p, 1e-10)); x = max(x, max(x) - 8); x = (x + 4) / 4; transpose.
    float32 throughout (the JS equivalent uses Float32Array + Math.log10).
    """
    p = np.asarray(power_mel, dtype=np.float32).reshape(FRAMES, N_MELS)
    x = np.log10(np.maximum(p, np.float32(1e-10))).astype(np.float32)
    x = np.maximum(x, x.max() - np.float32(8.0))
    x = ((x + np.float32(4.0)) / np.float32(4.0)).astype(np.float32)
    return np.ascontiguousarray(x.T)[None]


class HostTables:
    """Token-embedding lookup as the playground JS does it, plus the files that carry the tables.

    kind="int8": whisper_embed.i8.bin (int8 [51865,512], round(w / s)) + whisper_embed.scale.f32.bin
                 (s = rowwise absmax / 127): byte-identical to the previous export's files.
    kind="fp32": whisper_embed.f32.bin (fp32 [51865,512], the checkpoint's embed_tokens).
    Both: whisper_pos_embed.f32.bin (fp32 [448,512], decoder.embed_positions)."""

    def __init__(self, model, kind):
        e = model.model.decoder.embed_tokens.weight.detach().float()
        self.kind = kind
        self.pos = model.model.decoder.embed_positions.weight.detach().float().numpy()
        if kind == "int8":
            scale = e.abs().amax(dim=1) / 127.0
            self.table = torch.round(e / scale[:, None]).clamp(-127, 127).to(torch.int8).numpy()
            self.scale = scale.numpy().astype(np.float32)
        elif kind == "fp32":
            self.table = e.numpy().astype(np.float32)
            self.scale = None
        else:
            raise ValueError(kind)

    def embed(self, ids, position=None):
        e = self.table[ids].astype(np.float32)
        if self.scale is not None:
            e = e * self.scale[ids][..., None]
        if position is not None:
            e = e + self.pos[position]
        return e.astype(np.float32)

    def files(self):
        out = {}
        if self.kind == "int8":
            out["whisper_embed.i8.bin"] = self.table.tobytes()
            out["whisper_embed.scale.f32.bin"] = self.scale.astype("<f4").tobytes()
        else:
            out["whisper_embed.f32.bin"] = self.table.astype("<f4").tobytes()
        out["whisper_pos_embed.f32.bin"] = self.pos.astype("<f4").tobytes()
        return out


# --------------------------------------------------------------------------------------------
# Program wrappers.
# --------------------------------------------------------------------------------------------


def dft_mel_weights(feature_extractor):
    """cos / -sin DFT rows with the periodic Hann window folded in, plus the HF mel filters."""
    n = np.arange(N_FFT)
    k = np.arange(N_FREQ)
    window = 0.5 - 0.5 * np.cos(
        2.0 * np.pi * n / N_FFT
    )  # periodic Hann (torch / HF default)
    ang = 2.0 * np.pi * ((k[:, None] * n[None, :]) % N_FFT) / N_FFT
    wc = (np.cos(ang) * window[None, :]).astype(np.float32)
    ws = (-np.sin(ang) * window[None, :]).astype(np.float32)
    mel = np.asarray(feature_extractor.mel_filters, dtype=np.float64).T.astype(
        np.float32
    )
    assert mel.shape == (N_MELS, N_FREQ), mel.shape
    return torch.from_numpy(wc), torch.from_numpy(ws), torch.from_numpy(mel)


class MelProgram(torch.nn.Module):
    def __init__(self, feature_extractor):
        super().__init__()
        wc, ws, mel = dft_mel_weights(feature_extractor)
        self.register_buffer("wc", wc)
        self.register_buffer("ws", ws)
        self.register_buffer("mel", mel)

    def forward(self, frames):
        x = frames.reshape(
            FRAMES, N_FFT
        )  # aten.linear on the WebGPU runtime is 2-D only
        re = F.linear(x, self.wc)
        im = F.linear(x, self.ws)
        return F.linear(re * re + im * im, self.mel).reshape(1, FRAMES, N_MELS)


def _heads(x, s):
    """[s, 512] or [1, s, 512] -> [1, 8, s, 64]."""
    return x.reshape(1, s, N_HEADS, HEAD_DIM).transpose(1, 2)


def _merge(o, s):
    """[1, 8, s, 64] -> [1, s, 512]."""
    return o.transpose(1, 2).reshape(1, s, D_MODEL)


def attention(q, k, v, mask, impl):
    """q/k/v [1,H,S,D]; scale = 1/sqrt(64) = 0.125 (a power of two, so applying it to q as HF does
    or to q.k as the kernels do is bit-identical)."""
    if impl == "et_vk":
        import executorch.backends.vulkan.custom_ops_lib  # noqa: F401  (torch.ops.et_vk.*)

        return torch.ops.et_vk.sdpa(q, k, v, mask, None)
    s = torch.matmul(q * 0.125, k.transpose(-2, -1))
    if mask is not None:
        s = s + mask
    return torch.matmul(torch.softmax(s, dim=-1), v)


class EncoderProgram(torch.nn.Module):
    """HF WhisperEncoder forward re-expressed with 2-D fp32 linears.

    attn="eager": matmul + softmax (lowered to aten.bmm/_softmax; activation memory is planned).
    attn="et_vk": the fused et_vk.sdpa op (what the previous export used); on the public WebGPU
    runtime each call allocates its own two [8,1500,1500] f32 scratch buffers (72 MB each).
    """

    def __init__(self, model, attn="eager"):
        super().__init__()
        self.enc = model.model.encoder
        self.attn = attn

    def forward(self, log_mel):
        e = self.enc
        x = F.gelu(e.conv1(log_mel))
        x = F.gelu(e.conv2(x))  # [1, 512, 1500]
        x = x.permute(0, 2, 1).reshape(ENC_LEN, D_MODEL) + e.embed_positions.weight
        for layer in e.layers:
            a = layer.self_attn
            h = layer.self_attn_layer_norm(x)
            q = _heads(a.q_proj(h), ENC_LEN)
            k = _heads(a.k_proj(h), ENC_LEN)
            v = _heads(a.v_proj(h), ENC_LEN)
            o = attention(q, k, v, None, "et_vk" if self.attn == "et_vk" else "torch")
            x = x + a.out_proj(o.transpose(1, 2).reshape(ENC_LEN, D_MODEL))
            h = layer.final_layer_norm(x)
            x = x + layer.fc2(F.gelu(layer.fc1(h)))
        return e.layer_norm(x).reshape(1, ENC_LEN, D_MODEL)


def _lin(m, x):
    """nn.Linear that lowers on the public WebGPU runtime: fp32 aten.linear only accepts 2-D
    inputs, so fp32 weights get a [-1, K] view; int4 (q4gsw) linears accept any rank and keep it."""
    if type(m.weight.data) is torch.Tensor and x.dim() != 2:
        return m(x.reshape(-1, x.shape[-1])).reshape(*x.shape[:-1], m.out_features)
    return m(x)


class CrossKVProgram(torch.nn.Module):
    def __init__(self, decoder):
        super().__init__()
        self.layers = decoder.layers

    def forward(self, enc):
        out = []
        for layer in self.layers:
            a = layer.encoder_attn
            out.append(_heads(_lin(a.k_proj, enc), ENC_LEN))
            out.append(_heads(_lin(a.v_proj, enc), ENC_LEN))
        return tuple(out)


def _mlp(layer, x):
    return x + _lin(layer.fc2, F.gelu(_lin(layer.fc1, layer.final_layer_norm(x))))


def _cross(layer, x, k, v, s, impl):
    a = layer.encoder_attn
    h = layer.encoder_attn_layer_norm(x)
    o = attention(_heads(_lin(a.q_proj, h), s), k, v, None, impl)
    return x + _lin(a.out_proj, _merge(o, s))


class DecoderSplitProgram(torch.nn.Module):
    """Full-recompute decoder over MAX_DEC rows with a causal mask (no KV cache)."""

    def __init__(self, decoder, attn="et_vk"):
        super().__init__()
        self.dec = decoder
        self.attn = attn
        self.register_buffer(
            "pos", decoder.embed_positions.weight.detach()[:MAX_DEC].clone()
        )
        causal = torch.full((MAX_DEC, MAX_DEC), -1e30).triu(1)  # same fill as the previous export
        # et_vk.sdpa on the WebGPU runtime requires the full [B, H, S_q, S_kv] mask.
        self.register_buffer(
            "mask", causal.expand(1, N_HEADS, MAX_DEC, MAX_DEC).contiguous()
        )

    def forward(self, tok_embeds, *cross_kv):
        x = tok_embeds + self.pos
        for i, layer in enumerate(self.dec.layers):
            a = layer.self_attn
            h = layer.self_attn_layer_norm(x)
            q = _heads(_lin(a.q_proj, h), MAX_DEC)
            k = _heads(_lin(a.k_proj, h), MAX_DEC)
            v = _heads(_lin(a.v_proj, h), MAX_DEC)
            o = attention(q, k, v, self.mask, self.attn)
            x = x + _lin(a.out_proj, _merge(o, MAX_DEC))
            x = _cross(
                layer, x, cross_kv[2 * i], cross_kv[2 * i + 1], MAX_DEC, self.attn
            )
            x = _mlp(layer, x)
        return self.dec.layer_norm(x)


class DecodeStreamProgram(torch.nn.Module):
    """One token per call; self-attention K/V cache [1,448,8,64] per layer kept in the program.

    Self-attention uses llama.update_cache + llama.custom_sdpa (executorch.extension.llm.custom_ops),
    which the Vulkan partitioner fuses into sdpa_with_kv_cache with an et_vk.select_as_symint
    position. Cross-attention uses et_vk.sdpa (attn="et_vk") or matmul/softmax (attn="torch").
    """

    def __init__(self, decoder, attn="et_vk"):
        super().__init__()
        from executorch.extension.llm.custom_ops import custom_ops  # noqa: F401

        self.dec = decoder
        self.attn = attn
        for i in range(N_LAYERS):
            self.register_buffer(
                f"k_cache_{i}",
                torch.zeros(1, KV_LEN, N_HEADS, HEAD_DIM),
                persistent=False,
            )
            self.register_buffer(
                f"v_cache_{i}",
                torch.zeros(1, KV_LEN, N_HEADS, HEAD_DIM),
                persistent=False,
            )

    def reset(self):
        for i in range(N_LAYERS):
            getattr(self, f"k_cache_{i}").zero_()
            getattr(self, f"v_cache_{i}").zero_()

    def forward(self, embed, input_pos, *cross_kv):
        pos = input_pos[0].item()
        torch._check_is_size(pos)
        torch._check(pos < KV_LEN)
        x = embed
        for i, layer in enumerate(self.dec.layers):
            a = layer.self_attn
            h = layer.self_attn_layer_norm(x)
            q = _lin(a.q_proj, h).view(1, 1, N_HEADS, HEAD_DIM)
            k = _lin(a.k_proj, h).view(1, 1, N_HEADS, HEAD_DIM)
            v = _lin(a.v_proj, h).view(1, 1, N_HEADS, HEAD_DIM)
            kc = getattr(self, f"k_cache_{i}")
            vc = getattr(self, f"v_cache_{i}")
            torch.ops.llama.update_cache(k, kc, pos)
            torch.ops.llama.update_cache(v, vc, pos)
            o = torch.ops.llama.custom_sdpa(q, kc, vc, pos, None, 0.0, True, None)
            x = x + _lin(a.out_proj, o.reshape(1, 1, D_MODEL))
            x = _cross(layer, x, cross_kv[2 * i], cross_kv[2 * i + 1], 1, self.attn)
            x = _mlp(layer, x)
        return self.dec.layer_norm(x)


class LMHeadProgram(torch.nn.Module):
    """Tied output projection. The vocab is padded to `pad_to` rows: executorch 1.5.1's q4gsw
    lowering pads the scales to a multiple of 4 output channels and then fails its own shape
    check for N = 51865, so N must already be a multiple of 4 (the previous export used 51872).
    Padded rows have zero weights and a -1e9 bias, so they can never win an argmax."""

    def __init__(self, model, pad_to=0):
        super().__init__()
        w = model.model.decoder.embed_tokens.weight.detach().clone()
        n = w.shape[0]
        b = torch.zeros(max(n, pad_to))
        if pad_to and pad_to > n:
            w = torch.cat([w, w.new_zeros(pad_to - n, w.shape[1])])
            b[n:] = PAD_LOGIT
        self.proj = torch.nn.Linear(D_MODEL, w.shape[0], bias=True)
        self.proj.weight = torch.nn.Parameter(w, requires_grad=False)
        self.proj.bias = torch.nn.Parameter(b, requires_grad=False)

    def forward(self, hidden):
        return _lin(self.proj, hidden)


# --------------------------------------------------------------------------------------------
# Quantization.
# --------------------------------------------------------------------------------------------

Q4_DESC = (
    "torchao.quantization.quant_api.IntxWeightOnlyConfig(weight_dtype=torch.int4, "
    "granularity=PerGroup(64), intx_choose_qparams_algorithm='affine') applied with quantize_ to "
    "every nn.Linear of the program; symmetric, scale = group absmax / 7.5, q in [-8, 7]"
)


def quantize_q4(module, group_size=64):
    from torchao.quantization.granularity import PerGroup
    from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_

    quantize_(
        module,
        IntxWeightOnlyConfig(
            weight_dtype=torch.int4,
            granularity=PerGroup(group_size),
            intx_choose_qparams_algorithm="affine",
        ),
    )
    return module


Q8_HEAD_DESC = (
    "PT2E with executorch.backends.vulkan.quantizer.VulkanQuantizer: weights int8 per-output-channel "
    "symmetric (PerChannelMinMaxObserver, the quantizer's own weight spec), input activation int8 "
    "per-tensor affine STATIC (HistogramObserver) calibrated on lmhead inputs; lowers to "
    "quantized_decomposed.quantize_per_tensor + et_vk.linear_q8ta_q8csw"
)


def q8_lmhead(lm, calib_hidden):
    """Static-activation int8 x per-channel int8 lmhead (the only 8-bit linear executorch 1.5.1's
    Vulkan lowering emits that the public WebGPU runtime registers). Returns the converted GraphModule
    and the (input scale, zero point, dequantized weight) needed for an equivalent portable module."""
    from executorch.backends.vulkan.quantizer.vulkan_quantizer import (
        get_symmetric_quantization_config,
        VulkanQuantizer,
    )
    from torchao.quantization.pt2e import HistogramObserver
    from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
    from torchao.quantization.pt2e.quantizer import QuantizationConfig, QuantizationSpec

    base = get_symmetric_quantization_config(is_dynamic=False, weight_bits=8)
    act = QuantizationSpec(
        dtype=torch.int8,
        quant_min=-128,
        quant_max=127,
        qscheme=torch.per_tensor_affine,
        is_dynamic=False,
        observer_or_fake_quant_ctr=HistogramObserver,
    )
    cfg = QuantizationConfig(
        input_activation=act, output_activation=None, weight=base.weight, bias=None, is_qat=False
    )
    with torch.no_grad():
        gm = torch.export.export(lm.eval(), (calib_hidden[:1],), strict=False).module()
        gm = prepare_pt2e(gm, VulkanQuantizer().set_global(cfg))
        for i in range(calib_hidden.shape[0]):
            gm(calib_hidden[i : i + 1])
        gm = convert_pt2e(gm)
    qp, wdq = None, None
    for n in gm.graph.nodes:
        if n.op != "call_function":
            continue
        name = str(n.target)
        if "quantize_per_tensor" in name and "dequantize" not in name and qp is None:
            qp = (float(n.args[1]), int(n.args[2]))
        if "dequantize_per_channel" in name:
            w = getattr(gm, n.args[0].target)
            sc = getattr(gm, n.args[1].target)
            wdq = (w.float() * sc.float()[:, None]).contiguous()
    if qp is None or wdq is None:
        raise RuntimeError("q8 lmhead: could not find the quantize/dequantize nodes")
    return gm, qp, wdq


class FakeQ8LMHead(torch.nn.Module):
    """Portable-runtime equivalent of the q8 lmhead: per-tensor int8 fake-quant of the input, then
    an fp32 linear on the dequantized int8 weights (same bias)."""

    def __init__(self, lm, qp, wdq):
        super().__init__()
        self.scale, self.zp = qp
        self.register_buffer("w", wdq)
        self.register_buffer("b", lm.proj.bias.detach().clone())

    def forward(self, hidden):
        x = hidden.reshape(-1, D_MODEL)
        q = torch.clamp(torch.round(x / self.scale) + self.zp, -128, 127)
        x = (q - self.zp) * self.scale
        return F.linear(x, self.w, self.b).reshape(*hidden.shape[:-1], self.w.shape[0])


def word_errors(tokenizer, hyp_text, ref_text):
    """(edits, reference words) after WhisperTokenizer.normalize (English text normalizer)."""
    h = tokenizer.normalize(hyp_text).split()
    r = tokenizer.normalize(ref_text).split()
    d = list(range(len(r) + 1))
    for i in range(1, len(h) + 1):
        prev, d[0] = d[0], i
        for j in range(1, len(r) + 1):
            cur = d[j]
            d[j] = min(d[j] + 1, d[j - 1] + 1, prev + (h[i - 1] != r[j - 1]))
            prev = cur
    return d[len(r)], len(r)


def dequantized_copy(module):
    """Same module with every quantized linear weight replaced by its fp32 dequantization.

    Used only for the CPU portable .pte (the portable runtime has no torchao kernels); its numerics
    are the q4 numerics (dequantize-then-fp32-matmul, as the q4gsw kernels compute)."""
    m = copy.deepcopy(module)
    for sub in m.modules():
        if isinstance(sub, torch.nn.Linear) and hasattr(sub.weight, "dequantize"):
            sub.weight = torch.nn.Parameter(
                sub.weight.dequantize().float(), requires_grad=False
            )
    return m


# --------------------------------------------------------------------------------------------
# Speech sample.
# --------------------------------------------------------------------------------------------


def load_librispeech_dummy(parquet_path, max_sec=15.0, count=8, longer=False):
    """First `count` clips (in file order) no longer than max_sec seconds, as 16 kHz float32.
    longer=True selects the clips *longer* than max_sec instead (used only for calibration, so
    calibration and evaluation clips never overlap)."""
    import pyarrow.parquet as pq
    import soundfile as sf

    table = pq.read_table(parquet_path).to_pylist()
    clips = []
    for row in table:
        audio, sr = sf.read(io.BytesIO(row["audio"]["bytes"]), dtype="float32")
        if sr != SAMPLE_RATE:
            raise ValueError(f"unexpected sample rate {sr}")
        if audio.ndim != 1:
            audio = audio.mean(axis=1)
        if (audio.size / SAMPLE_RATE > max_sec) != longer:
            continue
        clips.append({"id": row["id"], "text": row["text"], "pcm": audio})
        if len(clips) >= count:
            break
    return clips


def seconds(pcm):
    return round(len(pcm) / SAMPLE_RATE, 3)


def isclose_report(a, b):
    a = np.asarray(a, dtype=np.float64).ravel()
    b = np.asarray(b, dtype=np.float64).ravel()
    d = np.abs(a - b)
    cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))
    return {
        "max_abs": float(d.max()),
        "mean_abs": float(d.mean()),
        "max_rel_to_absmax": float(d.max() / (np.abs(b).max() + 1e-30)),
        "cosine": cos,
    }



# --------------------------------------------------------------------------------------------
# Lowered-graph inspection (what the WebGPU runtime will see).
# --------------------------------------------------------------------------------------------


def _vk_graphs(pte_bytes):
    from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
        extract_vk_flatbuffer,
        flatbuffer_to_vk_graph,
    )
    from executorch.exir._serialize._program import deserialize_pte_binary

    pte = deserialize_pte_binary(pte_bytes)
    prog = pte.program if hasattr(pte, "program") else pte
    seg_base = int.from_bytes(pte_bytes[24:32], "little")
    out = []
    for plan in prog.execution_plan:
        for d in plan.delegates:
            proc = d.processed
            if proc.location.name == "INLINE":
                blob = prog.backend_delegate_data[proc.index].data
            else:
                seg = prog.segments[proc.index]
                blob = pte_bytes[seg_base + seg.offset : seg_base + seg.offset + seg.size]
            out.append(flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob)))
    return prog, out


def program_io(pte_bytes):
    """Method names and their input/output tensor shapes + dtypes from the serialized program."""
    from executorch.exir._serialize._program import deserialize_pte_binary

    pte = deserialize_pte_binary(pte_bytes)
    prog = pte.program if hasattr(pte, "program") else pte
    res = {}
    for plan in prog.execution_plan:

        def desc(i):
            v = plan.values[i].val
            if hasattr(v, "sizes"):
                return {"shape": list(v.sizes), "dtype": v.scalar_type.name.lower()}
            return {"kind": type(v).__name__}

        res[plan.name] = {
            "inputs": [desc(i) for i in plan.inputs],
            "outputs": [desc(i) for i in plan.outputs],
        }
    return res


def vk_checks(pte_bytes, runtime_ops=None):
    """Checks against constraints of the public WebGPU runtime's op implementations
    (backends/webgpu/runtime/ops): op registered, aten.linear 2-D, et_vk.sdpa mask shape,
    sdpa_with_kv_cache without mask / with causal, fp32 tensors only."""
    _, graphs = _vk_graphs(pte_bytes)
    problems, notes = [], {}
    for g in graphs:
        vals = [v.value for v in g.values]

        def t(i):
            v = vals[i]
            return v if type(v).__name__ == "VkTensor" else None

        dtypes = sorted({type(v).__name__ if type(v).__name__ != "VkTensor" else v.datatype.name for v in vals})
        notes["value_kinds"] = dtypes
        for op in g.chain:
            name = op.name
            if runtime_ops is not None and name not in runtime_ops:
                problems.append(f"{name}: not registered in the WebGPU runtime op list")
            a = op.args
            if name == "aten.linear.default":
                ri, rw = len(t(a[0]).dims), len(t(a[1]).dims)
                notes.setdefault("aten.linear input ranks", set()).add(ri)
                if ri != 2 or rw != 2:
                    problems.append(f"aten.linear with input rank {ri}, weight rank {rw} (runtime needs 2-D)")
            if name == "et_vk.sdpa.default":
                q, k = t(a[0]), t(a[1])
                m = t(a[3])
                if m is not None:
                    need = q.dims[0] * q.dims[1] * q.dims[2] * k.dims[2]
                    got = int(np.prod(m.dims))
                    if need != got:
                        problems.append(f"et_vk.sdpa mask numel {got} != B*H*Sq*Skv {need}")
                notes.setdefault("et_vk.sdpa q/k/mask", set()).add(
                    (tuple(q.dims), tuple(k.dims), tuple(m.dims) if m is not None else None)
                )
            if name == "sdpa_with_kv_cache.default":
                mask_kind = type(vals[a[7]]).__name__
                causal = vals[a[9]]
                if mask_kind != "Null":
                    problems.append("sdpa_with_kv_cache with attn_mask (runtime rejects)")
                if not (type(causal).__name__ == "Bool" and causal.bool_val):
                    problems.append(f"sdpa_with_kv_cache is_causal={causal} (runtime needs true)")
                notes.setdefault("sdpa_with_kv_cache q/k_cache", set()).add(
                    (tuple(t(a[0]).dims), tuple(t(a[3]).dims), type(vals[a[5]]).__name__)
                )
            if name == "aten.convolution.default":
                notes.setdefault("aten.convolution input/weight", set()).add(
                    (tuple(t(a[0]).dims), tuple(t(a[1]).dims))
                )
        for v in vals:
            if type(v).__name__ == "VkTensor" and v.datatype.name not in ("FLOAT32", "INT32", "INT64", "UINT8", "INT8"):
                problems.append(f"tensor dtype {v.datatype.name}")
                break
    for k, v in list(notes.items()):
        if isinstance(v, set):
            notes[k] = sorted(str(x) for x in v)
    return {"problems": problems, "ok": not problems, "notes": notes}