lumbda/examples/cuda-fanout/bench_binary.py

"""bench_binary.py — binary portal format vs text portal.

Same workloads as bench.py, but uses the new binary input/output
format that skips hex parsing entirely. This is where GPU should
actually start winning at heavy workloads."""

import hashlib
import os
import secrets
import struct
import subprocess
import sys
import time

BIN = sys.argv[1] if len(sys.argv) > 1 else "./shake256-fanout"
OUT_BYTES = 32


def write_binary_input(path, inputs_bytes, out_bytes):
    """Format: u32 out_bytes | u32 n | (u32 len | len bytes) × n"""
    with open(path, "wb") as fh:
        fh.write(struct.pack("<II", out_bytes, len(inputs_bytes)))
        for blob in inputs_bytes:
            fh.write(struct.pack("<I", len(blob)))
            fh.write(blob)


def read_binary_output(path):
    """Format: u32 n | u32 out_bytes | (out_bytes bytes) × n"""
    with open(path, "rb") as fh:
        n, out_bytes = struct.unpack("<II", fh.read(8))
        return [fh.read(out_bytes) for _ in range(n)]


def host_bench(inputs_bytes, out_bytes):
    t0 = time.monotonic()
    out = [hashlib.shake_256(b).digest(out_bytes) for b in inputs_bytes]
    return out, time.monotonic() - t0


def device_bench(inputs_bytes, out_bytes, use_daemon_proc=None):
    in_p = "/tmp/cf-bin-in.bin"
    out_p = "/tmp/cf-bin-out.bin"
    write_binary_input(in_p, inputs_bytes, out_bytes)
    t0 = time.monotonic()
    if use_daemon_proc is not None:
        use_daemon_proc.stdin.write(f"process-bin {in_p} {out_p}\n")
        use_daemon_proc.stdin.flush()
        line = use_daemon_proc.stdout.readline().strip()
        if not line.startswith("done"):
            raise RuntimeError(f"daemon error: {line}")
    else:
        subprocess.run([BIN, "--binary", in_p, out_p], check=True)
    wall = time.monotonic() - t0
    out = read_binary_output(out_p)
    os.unlink(in_p); os.unlink(out_p)
    return out, wall


GRID = [
    # (input_size_bytes, N)
    (32,            10_000),
    (32,           100_000),
    (32,         1_000_000),
    (1024,           1_000),
    (1024,          10_000),
    (1024,         100_000),
    (16 * 1024,      1_000),
    (16 * 1024,     10_000),
    (256 * 1024,       100),
    (256 * 1024,     1_000),
]


def main():
    print(f"binary: {BIN}")
    print(f"out_bytes: {OUT_BYTES}")
    print()

    # Start a daemon once so we exclude the 200ms init from the
    # interesting per-call columns.
    proc = subprocess.Popen([BIN, "--daemon"],
                            stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
    ready = proc.stdout.readline().strip()
    assert ready == "ready", f"daemon ready={ready!r}"

    print(f"{'in_sz':>10}  {'N':>9}  {'total':>9}  {'host ms':>9}  "
          f"{'dev (1shot)':>11}  {'dev (daem)':>10}  {'spd_daem':>9}  match")
    print("-" * 94)
    sweet = []
    for sz, N in GRID:
        inputs_bytes = [secrets.token_bytes(sz) for _ in range(N)]
        host_out, host_t = host_bench(inputs_bytes, OUT_BYTES)
        # daemon
        dev_d_out, dev_d_t = device_bench(inputs_bytes, OUT_BYTES, use_daemon_proc=proc)
        # one-shot
        dev_1_out, dev_1_t = device_bench(inputs_bytes, OUT_BYTES)
        match = host_out == dev_d_out == dev_1_out
        speedup_d = host_t / dev_d_t if dev_d_t else float("inf")
        total = sz * N
        total_str = f"{total/1e6:.1f}MB" if total >= 1e6 else f"{total/1e3:.0f}KB"
        print(f"{sz:>10}  {N:>9,}  {total_str:>9}  {host_t*1e3:>9.0f}  "
              f"{dev_1_t*1e3:>11.0f}  {dev_d_t*1e3:>10.0f}  {speedup_d:>8.2f}x  {match}")
        if speedup_d > 1.0:
            sweet.append((sz, N, total, host_t, dev_d_t))

    proc.stdin.write("quit\n"); proc.stdin.flush()
    proc.stdout.readline(); proc.wait()

    print()
    if sweet:
        print("=== GPU wins (daemon mode, binary format) ===")
        for sz, N, tot, ht, dt in sweet:
            print(f"  input={sz:>8} × N={N:>7,} ({tot/1e6:>6.1f} MB):"
                  f"  host {ht*1e3:>6.0f} ms → device {dt*1e3:>6.0f} ms ({ht/dt:.2f}×)")
    else:
        print("=== GPU still loses everywhere ===")
    print()
    print("Columns: dev (1shot) includes 200 ms cuda init each call.")
    print("         dev (daem) reuses warm daemon — the real production cost.")


if __name__ == "__main__":
    main()