70cd436eb9
Extracted from lato.2_mrp_mlx. The same sparse module underlies LATO.2, Pixal3D and the rest of the TRELLIS.2 family, and all of them are blocked on Apple Silicon by the same single op - submanifold conv - so it belongs in one tested package rather than vendored per port. 13/13 tests: 7 for the conv against a hand-written reference (spconv is uninstallable here so there is no upstream oracle), 6 for the remaining layers against torch. SubMConv3d runs 13.8ms at 128^3/128ch on m3ultra, 5.9x faster than the obvious per-offset loop.
129 lines
4.2 KiB
Python
129 lines
4.2 KiB
Python
"""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 trellis_sparse_mlx.conv import SubMConv3d, build_indice_map # noqa: E402
|
|
from trellis_sparse_mlx.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()
|