| """Shared helpers for the ExecuTorch WebGPU export recipes in this directory. |
| |
| Each recipe builds an eager model from a pinned upstream revision and calls |
| `export_vision` (or its own LLM path), which writes: |
| |
| <out>/<name>_webgpu.pte program lowered with WebGPUPartitioner (delegate id "VulkanBackend") |
| <out>/<name>_portable.pte same exported graph without a delegate, for CPU validation only |
| <out>/recipe.json upstream id + revision, versions, shapes, dtype, command line |
| <out>/validation.json eager vs exported-program vs portable .pte (and goldens, if given) |
| """ |
|
|
| import hashlib |
| import io |
| import json |
| import platform |
| import shlex |
| import sys |
| import urllib.request |
| import zipfile |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| PROXY_HINT = ( |
| "set https_proxy if the machine needs a proxy for huggingface.co / pypi.org" |
| ) |
|
|
|
|
| def versions(): |
| import executorch |
| import torchao |
| from executorch.version import git_version |
|
|
| out = { |
| "python": platform.python_version(), |
| "torch": torch.__version__, |
| "torchao": torchao.__version__, |
| "executorch": getattr(executorch, "__version__", None) |
| or __import__("executorch.version").version.__version__, |
| "executorch_git": git_version, |
| } |
| try: |
| import transformers |
|
|
| out["transformers"] = transformers.__version__ |
| except ImportError: |
| pass |
| try: |
| import huggingface_hub |
|
|
| out["huggingface_hub"] = huggingface_hub.__version__ |
| except ImportError: |
| pass |
| return out |
|
|
|
|
| def sha256_file(path): |
| h = hashlib.sha256() |
| with open(path, "rb") as fh: |
| for chunk in iter(lambda: fh.read(1 << 20), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def hf_file(repo_id, filename, revision, expected_sha256=None, cache_dir=None): |
| from huggingface_hub import hf_hub_download |
|
|
| path = hf_hub_download( |
| repo_id=repo_id, filename=filename, revision=revision, cache_dir=cache_dir |
| ) |
| digest = sha256_file(path) |
| if expected_sha256 and digest != expected_sha256: |
| sys.exit( |
| f"{repo_id}@{revision}:{filename} sha256 {digest} != {expected_sha256}" |
| ) |
| return Path(path), digest |
|
|
|
|
| def pypi_wheel_member(url, expected_sha256, member): |
| """Read one file out of a PyPI wheel, verifying the wheel's sha256.""" |
| with urllib.request.urlopen(url, timeout=120) as resp: |
| blob = resp.read() |
| digest = hashlib.sha256(blob).hexdigest() |
| if digest != expected_sha256: |
| sys.exit(f"{url}: sha256 {digest} != {expected_sha256} ({PROXY_HINT})") |
| with zipfile.ZipFile(io.BytesIO(blob)) as zf: |
| return zf.read(member).decode() |
|
|
|
|
| def scrub(text, out_dir): |
| """Replace machine-specific paths so recipe.json can ship with the model.""" |
| text = text.replace(str(Path(out_dir).resolve()), "<out>") |
| text = text.replace(str(Path(__file__).resolve().parent), "recipes") |
| return text.replace(str(Path.home()), "~") |
|
|
|
|
| def command_line(out_dir): |
| argv = [f"recipes/{Path(sys.argv[0]).name}", *sys.argv[1:]] |
| args = [] |
| for a in argv: |
| if a.startswith("/") or a.startswith("."): |
| a = scrub(str(Path(a).resolve()), out_dir) |
| args.append(a) |
| return scrub(" ".join(["python", *(shlex.quote(a) for a in args)]), out_dir) |
|
|
|
|
| def lower_webgpu(ep): |
| from executorch.backends.webgpu.partitioner import WebGPUPartitioner |
| from executorch.exir import to_edge_transform_and_lower |
|
|
| return to_edge_transform_and_lower( |
| ep, partitioner=[WebGPUPartitioner()] |
| ).to_executorch() |
|
|
|
|
| def lower_portable(ep): |
| from executorch.exir import to_edge_transform_and_lower |
|
|
| return to_edge_transform_and_lower(ep).to_executorch() |
|
|
|
|
| def delegate_summary(pte_bytes): |
| """Non-delegated kernels and per-delegate op histograms of a serialized program.""" |
| import collections |
|
|
| 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") |
| plan = prog.execution_plan[0] |
| ops = [f"{o.name}.{o.overload}" for o in plan.operators] |
| kernels = collections.Counter() |
| for chain in plan.chains: |
| for ins in chain.instructions: |
| if type(ins.instr_args).__name__ == "KernelCall": |
| kernels[ops[ins.instr_args.op_index]] += 1 |
| delegates = [] |
| 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] |
| g = flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob)) |
| delegates.append( |
| { |
| "id": d.id, |
| "ops": dict(collections.Counter(c.name for c in g.chain)), |
| "inputs": len(g.input_ids), |
| "outputs": len(g.output_ids), |
| } |
| ) |
| return {"non_delegated_kernels": dict(kernels), "delegates": delegates} |
|
|
|
|
| def compare(a, b): |
| a = np.asarray(a, dtype=np.float64).ravel() |
| b = np.asarray(b, dtype=np.float64).ravel() |
| n = min(a.size, b.size) |
| a, b = a[:n], b[:n] |
| cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30)) |
| return { |
| "max_abs": float(np.max(np.abs(a - b))), |
| "mean_abs": float(np.mean(np.abs(a - b))), |
| "cosine": cos, |
| } |
|
|
|
|
| def run_pte(pte_bytes, inputs): |
| from executorch.extension.pybindings.portable_lib import ( |
| _load_for_executorch_from_buffer, |
| ) |
|
|
| m = _load_for_executorch_from_buffer(pte_bytes) |
| return [ |
| o.detach().numpy() if isinstance(o, torch.Tensor) else o |
| for o in m.forward(inputs) |
| ] |
|
|
|
|
| def try_run_webgpu_pte_on_cpu(pte_bytes, inputs): |
| try: |
| return run_pte(pte_bytes, inputs), None |
| except Exception as e: |
| return None, f"{type(e).__name__}: {str(e).splitlines()[0][:300]}" |
|
|
|
|
| def export_vision( |
| *, |
| name, |
| model, |
| example_input, |
| out_dir, |
| upstream, |
| extra_recipe, |
| golden=None, |
| output_index=None, |
| ): |
| """Export, lower, write both .pte files, recipe.json, and validation.json. |
| |
| `golden` is an optional (input ndarray, output ndarray) pair; when its input |
| shape matches `example_input`, all paths are also compared against it. |
| `output_index` picks one tensor when the eager model returns a tuple. |
| """ |
| out_dir = Path(out_dir) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| torch.manual_seed(0) |
| model = model.eval() |
|
|
| class Wrapped(torch.nn.Module): |
| def __init__(self, m): |
| super().__init__() |
| self.m = m |
|
|
| def forward(self, x): |
| y = self.m(x) |
| return y[output_index] if output_index is not None else y |
|
|
| wrapped = Wrapped(model).eval() |
| with torch.no_grad(): |
| ep = torch.export.export(wrapped, (example_input,), strict=False) |
| webgpu = lower_webgpu(ep) |
| portable = lower_portable(ep) |
|
|
| webgpu_path = out_dir / f"{name}_webgpu.pte" |
| portable_path = out_dir / f"{name}_portable.pte" |
| webgpu_path.write_bytes(webgpu.buffer) |
| portable_path.write_bytes(portable.buffer) |
|
|
| summary = delegate_summary(webgpu.buffer) |
| out_shape = list(ep.module()(example_input).shape) |
| recipe = { |
| "model": name, |
| "upstream": upstream, |
| "versions": versions(), |
| "export": { |
| "api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner()]) -> to_executorch()", |
| "partitioner": "executorch.backends.webgpu.partitioner.WebGPUPartitioner() (default compile options; wraps VulkanPartitioner)", |
| "inputs": [ |
| { |
| "shape": list(example_input.shape), |
| "dtype": "float32", |
| "layout": "NCHW", |
| } |
| ], |
| "outputs": [{"shape": out_shape, "dtype": "float32"}], |
| "dtype": "float32", |
| "quantization": "none (fp32)", |
| **extra_recipe, |
| }, |
| "files": { |
| webgpu_path.name: { |
| "bytes": webgpu_path.stat().st_size, |
| "sha256": sha256_file(webgpu_path), |
| }, |
| portable_path.name: { |
| "bytes": portable_path.stat().st_size, |
| "sha256": sha256_file(portable_path), |
| "note": "CPU-only validation artifact (portable kernels, no delegate); not for release", |
| }, |
| }, |
| "graph": summary, |
| "command": command_line(out_dir), |
| } |
| (out_dir / "recipe.json").write_text(json.dumps(recipe, indent=2) + "\n") |
|
|
| |
| gen = torch.Generator().manual_seed(1234) |
| rand_input = torch.rand(example_input.shape, generator=gen) |
| cases = {"random": rand_input} |
| if golden is not None and tuple(golden[0].shape) == tuple(example_input.shape): |
| cases["golden"] = torch.from_numpy(golden[0].copy()) |
| report = { |
| "cpu_runtime": "executorch pybindings (portable_lib) from the pip wheel", |
| "cases": {}, |
| } |
| webgpu_cpu_error = None |
| for label, x in cases.items(): |
| with torch.no_grad(): |
| eager = wrapped(x).numpy() |
| exported = ep.module()(x).numpy() |
| port = run_pte(portable.buffer, (x,))[0] |
| res = { |
| "exported_vs_eager": compare(exported, eager), |
| "portable_pte_vs_eager": compare(port, eager), |
| } |
| wg, err = try_run_webgpu_pte_on_cpu(webgpu.buffer, (x,)) |
| if wg is not None: |
| res["webgpu_pte_vs_eager"] = compare(wg[0], eager) |
| else: |
| webgpu_cpu_error = err |
| if label == "golden": |
| ref = golden[1] |
| res["internal_golden_vs_eager"] = compare(eager, ref) |
| res["internal_golden_vs_portable_pte"] = compare(port, ref) |
| res["golden_note"] = ( |
| "reference output from the internal goldens.manifest.json" |
| ) |
| report["cases"][label] = res |
| report["webgpu_pte_on_cpu"] = ( |
| "ran" |
| if webgpu_cpu_error is None |
| else f"not runnable on CPU pybindings: {webgpu_cpu_error}" |
| ) |
| (out_dir / "validation.json").write_text(json.dumps(report, indent=2) + "\n") |
| print( |
| json.dumps( |
| {"recipe": recipe["files"], "graph": summary, "validation": report}, |
| indent=2, |
| ) |
| ) |
| return ep, webgpu, portable |
|
|
|
|
| def load_golden(manifest_path, key): |
| man = json.loads(Path(manifest_path).read_text())[key]["tensors"] |
| base = Path(manifest_path).parent |
|
|
| def ld(t): |
| return np.fromfile(base / t["file"], dtype=np.float32).reshape(t["shape"]) |
|
|
| inp = man.get("input") |
| return ld(inp), ld(man["output"]) |
|
|
