lato.2_mrp_mlx/bench/bench_subm.py

"""Benchmark SubMConv3d across the fleet. Needs only mlx + numpy.

Submanifold conv is the hot op in the V-VAE, so its scaling across Apple Silicon
generations is what decides which box should run LATO.2 once the port lands. Two
distinct costs are timed separately because they scale differently:

  indice map — coordinate hashing + binary search, runs on CPU in numpy, cacheable
               per coordinate set (upstream reuses it via `indice_key`)
  conv       — 27 gathers + 27 matmuls on the GPU, paid every layer

Occupancy is set to ~10% of the grid, roughly what a voxelised surface gives you.
"""

from __future__ import annotations

import argparse
import json
import platform
import subprocess
import sys
import time
from pathlib import Path

import mlx.core as mx
import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from lato_mlx.sparse.conv import SubMConv3d, build_indice_map  # noqa: E402
from lato_mlx.sparse.tensor import SparseTensor  # noqa: E402


def machine() -> dict:
    def sysctl(k):
        try:
            return subprocess.run(
                ["sysctl", "-n", k], capture_output=True, text=True, timeout=5
            ).stdout.strip()
        except Exception:  # noqa: BLE001
            return "?"

    mem = sysctl("hw.memsize")
    return {
        "host": platform.node().split(".")[0],
        "chip": sysctl("machdep.cpu.brand_string"),
        "ram_gb": round(int(mem) / 1e9) if mem.isdigit() else "?",
        "mlx": mx.__version__,
    }


def make_sparse(res: int, channels: int, occupancy: float, seed: int = 0):
    rng = np.random.default_rng(seed)
    n_total = res**3
    n = max(64, int(n_total * occupancy))
    flat = rng.choice(n_total, size=n, replace=False)
    flat.sort()
    z, rem = np.divmod(flat, res * res)
    y, x = np.divmod(rem, res)
    coords = np.stack([np.zeros_like(z), z, y, x], axis=1).astype(np.int32)
    feats = rng.standard_normal((n, channels)).astype(np.float32)
    return mx.array(coords), mx.array(feats)


def time_it(fn, warmup=1, iters=3):
    for _ in range(warmup):
        fn()
    ts = []
    for _ in range(iters):
        t0 = time.perf_counter()
        fn()
        ts.append(time.perf_counter() - t0)
    return min(ts)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=None, help="write JSON result here")
    ap.add_argument("--occupancy", type=float, default=0.10)
    ap.add_argument("--iters", type=int, default=3)
    a = ap.parse_args()

    info = machine()
    print(f"== {info['host']}  {info['chip']}  {info['ram_gb']}GB  mlx {info['mlx']}")
    print(f"{'res':>5} {'chan':>5} {'voxels':>9} {'imap_ms':>9} {'conv_ms':>9} {'Mvox/s':>8}")

    rows = []
    for res, ch in [(32, 64), (64, 64), (64, 128), (128, 128)]:
        try:
            coords, feats = make_sparse(res, ch, a.occupancy)
            n = coords.shape[0]
            conv = SubMConv3d(ch, ch, 3, bias=True, indice_key=f"b{res}")

            t_imap = time_it(lambda: build_indice_map(coords, 3), warmup=0, iters=1)

            x = SparseTensor(feats, coords)
            _ = conv(x)  # populate the indice cache so we time conv alone
            mx.eval(_.feats)

            def run():
                mx.eval(conv(x).feats)

            t_conv = time_it(run, iters=a.iters)
            rows.append(
                {
                    "res": res,
                    "channels": ch,
                    "voxels": n,
                    "imap_ms": t_imap * 1e3,
                    "conv_ms": t_conv * 1e3,
                    "mvox_s": n / t_conv / 1e6,
                }
            )
            print(
                f"{res:>5} {ch:>5} {n:>9,} {t_imap*1e3:>9.1f} {t_conv*1e3:>9.1f} "
                f"{n/t_conv/1e6:>8.2f}"
            )
        except Exception as e:  # noqa: BLE001
            print(f"{res:>5} {ch:>5}  FAILED: {type(e).__name__}: {e}")
            rows.append({"res": res, "channels": ch, "error": f"{type(e).__name__}: {e}"})

    result = {"machine": info, "occupancy": a.occupancy, "rows": rows}
    if a.out:
        Path(a.out).parent.mkdir(parents=True, exist_ok=True)
        Path(a.out).write_text(json.dumps(result, indent=2))
        print(f"\nwrote {a.out}")


if __name__ == "__main__":
    main()