rlx-web 0.2.18

#!/usr/bin/env python3
# RLX — versatile ML compiler + runtime.
# Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Build the rlx-web WebAssembly bundle and (optionally) serve the demo.

One command, all platforms (macOS / Linux / Windows):

    python3 crates/rlx-web/build.py            # CPU build -> web/pkg
    python3 crates/rlx-web/build.py --webgpu   # also bring up a WebGPU device
    python3 crates/rlx-web/build.py --serve     # build + serve the demo

The `wasm-bindgen` CLI must match the `wasm-bindgen` crate version. If it is
missing or mismatched, this script prints the exact `cargo install` command.
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import webbrowser
from pathlib import Path

CRATE = "rlx-web"
WASM_ARTIFACT = "rlx_web.wasm"  # cargo replaces '-' with '_'
OUT_NAME = "rlx_web"
TARGET = "wasm32-unknown-unknown"

HERE = Path(__file__).resolve().parent          # crates/rlx-web
WORKSPACE = HERE.parent.parent                  # repo root
WEB_DIR = HERE / "web"
PKG_DIR = WEB_DIR / "pkg"


def run(cmd: list[str], **kw) -> subprocess.CompletedProcess:
    print("·", " ".join(cmd))
    return subprocess.run(cmd, check=True, **kw)


# WebAssembly SIMD — opt-in, because it was measured and did nothing.
#
# The reasoning for enabling it is sound: wasm has no baseline vector unit the
# way x86_64 (SSE2) and aarch64 (NEON) do, so the plain-Rust kernels that
# auto-vectorise natively compile one lane at a time, and the GEMM here is
# already the pure-Rust fallback (`rlx-cpu`'s BLAS path is an unsatisfiable
# import on wasm). The measurement disagreed: a cluster rank training
# `mnist-cnn` in Chrome did 1,244 samples/s scalar and 1,227 with `+simd128` —
# no change, for a bundle that grew 5.7 MB → 8.6 MB.
#
# The actual bottleneck was a clamped `setTimeout(0)` in the round loop
# (`web/cluster-protocol.js`), worth 4× once fixed. Left as a flag rather than
# deleted: it should pay off once the kernels are written to vectorise, and the
# next person should not have to rediscover that it currently does not.
SIMD_FLAG = "-C target-feature=+simd128"


def wasm_bindgen_crate_version() -> str:
    """Read the locked wasm-bindgen version from `cargo metadata`."""
    out = subprocess.run(
        ["cargo", "metadata", "--format-version", "1"],
        cwd=WORKSPACE, check=True, capture_output=True, text=True,
    )
    meta = json.loads(out.stdout)
    for pkg in meta["packages"]:
        if pkg["name"] == "wasm-bindgen":
            return pkg["version"]
    raise SystemExit("wasm-bindgen not found in cargo metadata")


def ensure_wasm_bindgen(expected: str) -> str:
    exe = shutil.which("wasm-bindgen")
    install = f"cargo install wasm-bindgen-cli --version {expected}"
    if exe is None:
        raise SystemExit(
            f"wasm-bindgen CLI not found.\n  Install the matching version:\n    {install}"
        )
    got = subprocess.run([exe, "--version"], check=True, capture_output=True, text=True)
    ver = got.stdout.strip().split()[-1]
    if ver != expected:
        raise SystemExit(
            f"wasm-bindgen CLI {ver} != crate {expected} (they must match).\n"
            f"  Fix:\n    {install}"
        )
    return exe


def target_dir() -> Path:
    out = subprocess.run(
        ["cargo", "metadata", "--format-version", "1", "--no-deps"],
        cwd=WORKSPACE, check=True, capture_output=True, text=True,
    )
    return Path(json.loads(out.stdout)["target_directory"])


BACKENDS = ("python", "npx", "miniserve", "basic-http-server")


def serve_static(web_dir: Path, backend: str, port: int) -> None:
    """Serve `web_dir` with the chosen static file server."""
    if backend == "python":
        import functools
        import http.server
        import socketserver

        handler = functools.partial(http.server.SimpleHTTPRequestHandler, directory=str(web_dir))
        with socketserver.TCPServer(("", port), handler) as httpd:
            httpd.serve_forever()
        return

    if backend == "npx":
        if shutil.which("npx") is None:
            raise SystemExit("npx not found — use --serve-with python")
        subprocess.run(["npx", "--yes", "serve", str(web_dir), "-l", str(port)], check=True)
        return

    if backend == "miniserve":
        exe = shutil.which("miniserve") or shutil.which("miniserve.exe")
        if exe is None:
            raise SystemExit("miniserve not found — cargo install miniserve")
        subprocess.run([exe, "-p", str(port), str(web_dir)], check=True)
        return

    if backend == "basic-http-server":
        exe = shutil.which("basic-http-server")
        if exe is None:
            raise SystemExit("basic-http-server not found — cargo install basic-http-server")
        subprocess.run([exe, str(web_dir), "--addr", f"127.0.0.1:{port}"], check=True)
        return

    raise SystemExit(f"unknown serve backend {backend!r}")


def main() -> None:
    ap = argparse.ArgumentParser(description="Build the rlx-web wasm bundle.")
    ap.add_argument("--webgpu", action="store_true", help="enable the WebGPU compute path")
    ap.add_argument("--webgl", action="store_true", help="enable the WebGL2 GPGPU path")
    ap.add_argument("--all", action="store_true", help="enable every GPU backend")
    ap.add_argument("--serve", action="store_true", help="serve the demo after building")
    ap.add_argument(
        "--serve-with",
        choices=BACKENDS,
        default="python",
        help="static server for --serve (default: python stdlib)",
    )
    ap.add_argument(
        "--simd",
        action="store_true",
        help="build with wasm SIMD (simd128). Measured as no gain and +50% size; see SIMD_FLAG",
    )
    ap.add_argument("--debug", action="store_true", help="debug build (default: release)")
    ap.add_argument("--port", type=int, default=8000, help="port for --serve")
    args = ap.parse_args()

    expected = wasm_bindgen_crate_version()
    bindgen = ensure_wasm_bindgen(expected)

    run(["rustup", "target", "add", TARGET])

    features = []
    if args.webgpu or args.all:
        features.append("webgpu")
    if args.webgl or args.all:
        features.append("webgl")

    cargo = ["cargo", "build", "-p", CRATE, "--target", TARGET]
    if not args.debug:
        cargo.append("--release")
    if features:
        cargo += ["--features", ",".join(features)]

    env = dict(os.environ)
    if args.simd:
        existing = env.get("RUSTFLAGS", "").strip()
        env["RUSTFLAGS"] = f"{existing} {SIMD_FLAG}".strip()
        print(f"· RUSTFLAGS={env['RUSTFLAGS']}")
    run(cargo, cwd=WORKSPACE, env=env)

    profile = "debug" if args.debug else "release"
    wasm = target_dir() / TARGET / profile / WASM_ARTIFACT
    if not wasm.exists():
        raise SystemExit(f"expected wasm artifact not found: {wasm}")

    PKG_DIR.mkdir(parents=True, exist_ok=True)
    run([bindgen, "--target", "web", "--typescript", "--out-dir", str(PKG_DIR),
         "--out-name", OUT_NAME, str(wasm)])

    print(f"\n✓ bundle ready: {PKG_DIR}")
    print(f"  backends: {', '.join(['cpu', *features]) if features else 'cpu'}")

    if args.serve:
        url = f"http://127.0.0.1:{args.port}/vision-bench.html"
        print(f"\nServing {WEB_DIR} via {args.serve_with} at {url}  (Ctrl-C to stop)")
        print("  index.html — MLP demo")
        print("  vision-bench.html — vision models")
        try:
            webbrowser.open(url)
        except Exception:
            pass
        try:
            serve_static(WEB_DIR, args.serve_with, args.port)
        except KeyboardInterrupt:
            print("\nstopped.")


if __name__ == "__main__":
    main()