code/scripts/bench.py · changh95/hamer-p150 at main

p150 ETH-dispatch compliance (2026-10-05): default ETH dispatch, 1 CQ, 12x10 in Python API and server; numbers re-measured

843986e verified about 1 month ago

Raw History Blame Contribute Delete

10.1 kB

# SPDX-License-Identifier: Apache-2.0
"""Baseline perf breakdown of the served HaMeR path on one Blackhole chip.
Measures, on the REAL checkpoint and the media/ sample crops (fresh input every request,
no same-input hash cache):
  device_trace_ms   execute_trace back-to-back (non-blocking, one sync at the end) / N
                    = pure device time of the traced ViT + MANO head graph
  per-request stage timings of the served path (median of N):
      im2col        host input prep: F.unfold -> fp32 [1,1,192,768] (patch=fp32w), or just the
                    fp32 contiguous crop (patch=device: the im2col runs on device)
      h2d           ttnn.from_torch(host) + copy_host_to_device_tensor (590 KB fp32)
      exec          execute_trace(blocking=True)
      d2h           ttnn.to_torch(dec [1,1,1,109])
      finalize      host init-param add + rot6d -> rotmat
      total         sum (what one served request costs on the device side)
  PCC vs the torch CPU reference per sample (trace replay with fresh input).
Grid / dispatch selection (must be decided before ttnn creates its runtime options):
  --dispatch eth|worker   ETH dispatch (default, patched tree, 1 CQ) or Tensix-column dispatch (opt-in)
  --grid    XxY (<=12x10) cap the compute grid via TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE
                          (11x10 = p150 stock, 12x10 = p150 + ETH dispatch / Galaxy stock)
"""
from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import time
from pathlib import Path

_REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO))


def parse():
    p = argparse.ArgumentParser()
    p.add_argument("--dispatch", default="eth", choices=["eth", "worker"],
                   help="eth (default: p150-equivalent, 1 CQ, 12x10) | worker (Galaxy-only opt-in; 11x10 on a p150)")
    p.add_argument("--grid", default="12x10", help="12x10 (policy grid) or smaller, e.g. 11x10 = stock p150")
    p.add_argument("--iters", type=int, default=50)
    p.add_argument("--ckpt", default=None)
    p.add_argument("--mean", default=None)
    p.add_argument("--json", default=None, help="write results JSON here")
    return p.parse_args()


def main():
    a = parse()
    # POLICY: tune/report on the 12x10 grid only (13x10 Galaxy-only grid is forbidden).
    x, y = (int(v) for v in a.grid.lower().split("x"))
    if x > 12 or y > 10:
        raise SystemExit("grid larger than 12x10 is not allowed (p150 12x10 policy)")
    os.environ["TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE"] = f"{x - 1},{y - 1}"
    os.environ.setdefault("TT_FUSED_HASH", "0")  # served config: fresh image every request

    import torch
    import ttnn
    from hamer import Hamer, build_from_reference, load_hamer_checkpoint
    from hamer import model as hm
    from hamer import mano_head as mh
    from hamer.fused import host_im2col
    from hamer.preprocess import preprocess

    hf = Path.home() / ".cache/huggingface/hub/models--changh95--hamer-weights/snapshots/e8217b46d54b4e8afaaf0dd6efb6b62341c97d32"
    ckpt = Path(a.ckpt) if a.ckpt else hf / "hamer.ckpt"
    mean = Path(a.mean) if a.mean else hf / "mano_mean_params.npz"
    torch.manual_seed(0)
    ref = Hamer().eval()
    load_hamer_checkpoint(ref, ckpt, mean)

    # dispatch core type: the model's own HAMER_DISPATCH switch (worker | eth)
    os.environ["HAMER_DISPATCH"] = a.dispatch

    tt = build_from_reference(ref, device_id=0, strict=True)
    g = tt.device.compute_with_storage_grid_size()
    print(f"[bench] dispatch={tt.dispatch or a.dispatch} cqs=1 grid={g.x}x{g.y} fused={tt.fused_opts.describe()}")

    media = _REPO.parent / "media"
    boxes = json.loads((media / "bboxes.json").read_text())
    xs = []
    xs_raw = []  # uint8 crops: what the server hands a TT_FUSED_INPUT=u8 model
    for name in sorted(boxes):
        e = boxes[name]
        if isinstance(e, list) and e and isinstance(e[0], dict):
            bbox, right = e[0]["bbox"], bool(e[0].get("is_right", 1))
        elif isinstance(e, dict):
            bbox, right = e["bbox"], bool(e.get("is_right", 1))
        else:
            bbox, right = e, True
        x, _ = preprocess(media / name, bbox_xyxy=tuple(bbox), is_right=right)
        xr, _ = preprocess(media / name, bbox_xyxy=tuple(bbox), is_right=right, raw=True)
        xs.append((name, x))
        xs_raw.append(xr)

    def pcc(u, v):
        u = u.reshape(-1).float(); v = v.reshape(-1).float()
        u = u - u.mean(); v = v - v.mean()
        return float((u @ v) / (u.norm() * v.norm()).clamp_min(1e-12))

    with torch.no_grad():
        refs = [ref(x) for _, x in xs]
        tt(xs[0][1])  # eager (JIT)
        tt(xs[1][1])  # trace capture + replay
        assert tt.trace_active
        pccs = {n: pcc(tt(x), r) for (n, x), r in zip(xs, refs)}
    print("[bench] PCC (trace replay, fresh input):", {k: round(v, 6) for k, v in pccs.items()})

    trace_id, dec_buf = tt._trace
    dev = tt.device
    # pure device time: back-to-back replays
    for _ in range(5):
        ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
    ttnn.synchronize_device(dev)
    t0 = time.perf_counter()
    for _ in range(a.iters):
        ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
    ttnn.synchronize_device(dev)
    dev_ms = (time.perf_counter() - t0) / a.iters * 1e3
    # 10 chunks of back-to-back replays -> median / min of the chunk means (noise on the shared host)
    chunks = []
    for _ in range(10):
        t0 = time.perf_counter()
        for _ in range(20):
            ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
        ttnn.synchronize_device(dev)
        chunks.append((time.perf_counter() - t0) / 20 * 1e3)
    dev_med, dev_min = statistics.median(chunks), min(chunks)
    print(f"[bench] device_trace_ms={dev_ms:.3f} chunk median={dev_med:.3f} min={dev_min:.3f}")
    if not tt.fused_opts.device_patch_embed:  # stage loop below assumes the fused im2col input
        print("RESULT " + json.dumps({"dispatch": a.dispatch, "grid": f"{g.x}x{g.y}", "device_trace_ms": dev_ms, "pcc": pccs}))
        tt.close()
        return

    # served path, stage by stage, fresh input each request
    u8 = tt.fused_opts.input == "u8"  # served input = the uint8 crop (server: preprocess_pil(raw=True))
    hp = tt._head_params
    io = tt._io
    if io is not None:
        # round 13 fast host I/O (TtHamer._forward_fast, stage by stage): prep = centring into the persistent host
        # buffer, h2d = copy_host_to_device_tensor, exec = replay + readback into the persistent host tensor + the
        # wait (spin on a CQ event, or a blocking read with HAMER_SPIN_WAIT=0), finalize = HostFinalizer
        st = {k: [] for k in ("prep", "h2d", "exec", "finalize", "total")}
        for i in range(a.iters):
            x = xs_raw[i % len(xs)] if u8 else xs[i % len(xs)][1]
            c0 = time.perf_counter()
            tt._write_input(io, x)  # centring (round 14: + the tiled im2col with TT_FUSED_TILE_IN)
            c1 = time.perf_counter()
            ttnn.copy_host_to_device_tensor(io["host_in"], tt._tokens_dev)
            c2 = time.perf_counter()
            ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
            if tt.spin_wait:
                ttnn.copy_device_to_host_tensor(dec_buf, io["host_out"], blocking=False)
                ttnn.event_synchronize(ttnn.record_event(dev, 0))
            else:
                ttnn.copy_device_to_host_tensor(dec_buf, io["host_out"], blocking=True)
            c3 = time.perf_counter()
            out = io["fin"](io["out_np"])
            c4 = time.perf_counter()
            for k, (s_, e_) in zip(st, [(c0, c1), (c1, c2), (c2, c3), (c3, c4), (c0, c4)]):
                st[k].append((e_ - s_) * 1e3)
    else:
      st = {k: [] for k in ("im2col", "h2d", "exec", "d2h", "finalize", "total")}
      for i in range(a.iters):
        x = xs_raw[i % len(xs)] if u8 else xs[i % len(xs)][1]
        c0 = time.perf_counter()
        host, hdt, lay = tt._host_input(x)
        c1 = time.perf_counter()
        h = ttnn.from_torch(host, device=None, layout=lay, dtype=hdt,  # round 15: the input layout (TILE with TT_FUSED_TILE_IN)
                            mesh_mapper=ttnn.ReplicateTensorToMesh(dev))
        ttnn.copy_host_to_device_tensor(h, tt._tokens_dev)
        c2 = time.perf_counter()
        ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=True)
        c3 = time.perf_counter()
        dec_h = tt._dec_host(ttnn.to_torch(dec_buf), 1)
        c4 = time.perf_counter()
        out = mh.host_finalize(dec_h, hp["dec_split"], tt._init_pose, tt._init_betas, tt._init_cam, init=tt._fin_init)
        c5 = time.perf_counter()
        for k, (s, e) in zip(st, [(c0, c1), (c1, c2), (c2, c3), (c3, c4), (c4, c5), (c0, c5)]):
            st[k].append((e - s) * 1e3)
    med = {k: statistics.median(v) for k, v in st.items()}
    # last output sanity vs reference
    last = (a.iters - 1) % len(xs)
    med_pcc_last = pcc(out, refs[last])
    # full TtHamer.__call__ (what the server calls), fresh input
    call = []
    with torch.no_grad():
        for i in range(a.iters):
            x = xs_raw[i % len(xs)] if u8 else xs[i % len(xs)][1]
            c0 = time.perf_counter(); tt(x); call.append((time.perf_counter() - c0) * 1e3)
    res = {
        "dispatch": a.dispatch, "grid": f"{g.x}x{g.y}", "device_trace_ms": dev_ms,
        "device_trace_chunk_median_ms": dev_med, "device_trace_chunk_min_ms": dev_min,
        "stages_ms": med, "ttHamer_call_ms": statistics.median(call), "ttHamer_call_min_ms": min(call),
        "pcc": pccs, "pcc_bench_last": med_pcc_last,
    }
    print(f"[bench] device_trace_ms={dev_ms:.3f}  (med {dev_med:.3f} min {dev_min:.3f})  served-call={res['ttHamer_call_ms']:.3f} ms (min {res['ttHamer_call_min_ms']:.3f})")
    print("[bench] stages (median ms): " + "  ".join(f"{k}={v:.3f}" for k, v in med.items()))
    print("RESULT " + json.dumps(res))
    if a.json:
        Path(a.json).write_text(json.dumps(res, indent=2))
    tt.close()


if __name__ == "__main__":
    main()

	# SPDX-License-Identifier: Apache-2.0
	"""Baseline perf breakdown of the served HaMeR path on one Blackhole chip.

	Measures, on the REAL checkpoint and the media/ sample crops (fresh input every request,
	no same-input hash cache):

	device_trace_ms execute_trace back-to-back (non-blocking, one sync at the end) / N
	= pure device time of the traced ViT + MANO head graph
	per-request stage timings of the served path (median of N):
	im2col host input prep: F.unfold -> fp32 [1,1,192,768] (patch=fp32w), or just the
	fp32 contiguous crop (patch=device: the im2col runs on device)
	h2d ttnn.from_torch(host) + copy_host_to_device_tensor (590 KB fp32)
	exec execute_trace(blocking=True)
	d2h ttnn.to_torch(dec [1,1,1,109])
	finalize host init-param add + rot6d -> rotmat
	total sum (what one served request costs on the device side)
	PCC vs the torch CPU reference per sample (trace replay with fresh input).

	Grid / dispatch selection (must be decided before ttnn creates its runtime options):
	--dispatch eth\|worker ETH dispatch (default, patched tree, 1 CQ) or Tensix-column dispatch (opt-in)
	--grid XxY (<=12x10) cap the compute grid via TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE
	(11x10 = p150 stock, 12x10 = p150 + ETH dispatch / Galaxy stock)
	"""
	from __future__ import annotations

	import argparse
	import json
	import os
	import statistics
	import sys
	import time
	from pathlib import Path

	_REPO = Path(__file__).resolve().parents[1]
	sys.path.insert(0, str(_REPO))


	def parse():
	p = argparse.ArgumentParser()
	p.add_argument("--dispatch", default="eth", choices=["eth", "worker"],
	help="eth (default: p150-equivalent, 1 CQ, 12x10) \| worker (Galaxy-only opt-in; 11x10 on a p150)")
	p.add_argument("--grid", default="12x10", help="12x10 (policy grid) or smaller, e.g. 11x10 = stock p150")
	p.add_argument("--iters", type=int, default=50)
	p.add_argument("--ckpt", default=None)
	p.add_argument("--mean", default=None)
	p.add_argument("--json", default=None, help="write results JSON here")
	return p.parse_args()


	def main():
	a = parse()
	# POLICY: tune/report on the 12x10 grid only (13x10 Galaxy-only grid is forbidden).
	x, y = (int(v) for v in a.grid.lower().split("x"))
	if x > 12 or y > 10:
	raise SystemExit("grid larger than 12x10 is not allowed (p150 12x10 policy)")
	os.environ["TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE"] = f"{x - 1},{y - 1}"
	os.environ.setdefault("TT_FUSED_HASH", "0") # served config: fresh image every request

	import torch
	import ttnn
	from hamer import Hamer, build_from_reference, load_hamer_checkpoint
	from hamer import model as hm
	from hamer import mano_head as mh
	from hamer.fused import host_im2col
	from hamer.preprocess import preprocess

	hf = Path.home() / ".cache/huggingface/hub/models--changh95--hamer-weights/snapshots/e8217b46d54b4e8afaaf0dd6efb6b62341c97d32"
	ckpt = Path(a.ckpt) if a.ckpt else hf / "hamer.ckpt"
	mean = Path(a.mean) if a.mean else hf / "mano_mean_params.npz"
	torch.manual_seed(0)
	ref = Hamer().eval()
	load_hamer_checkpoint(ref, ckpt, mean)

	# dispatch core type: the model's own HAMER_DISPATCH switch (worker \| eth)
	os.environ["HAMER_DISPATCH"] = a.dispatch

	tt = build_from_reference(ref, device_id=0, strict=True)
	g = tt.device.compute_with_storage_grid_size()
	print(f"[bench] dispatch={tt.dispatch or a.dispatch} cqs=1 grid={g.x}x{g.y} fused={tt.fused_opts.describe()}")

	media = _REPO.parent / "media"
	boxes = json.loads((media / "bboxes.json").read_text())
	xs = []
	xs_raw = [] # uint8 crops: what the server hands a TT_FUSED_INPUT=u8 model
	for name in sorted(boxes):
	e = boxes[name]
	if isinstance(e, list) and e and isinstance(e[0], dict):
	bbox, right = e[0]["bbox"], bool(e[0].get("is_right", 1))
	elif isinstance(e, dict):
	bbox, right = e["bbox"], bool(e.get("is_right", 1))
	else:
	bbox, right = e, True
	x, _ = preprocess(media / name, bbox_xyxy=tuple(bbox), is_right=right)
	xr, _ = preprocess(media / name, bbox_xyxy=tuple(bbox), is_right=right, raw=True)
	xs.append((name, x))
	xs_raw.append(xr)

	def pcc(u, v):
	u = u.reshape(-1).float(); v = v.reshape(-1).float()
	u = u - u.mean(); v = v - v.mean()
	return float((u @ v) / (u.norm() * v.norm()).clamp_min(1e-12))

	with torch.no_grad():
	refs = [ref(x) for _, x in xs]
	tt(xs[0][1]) # eager (JIT)
	tt(xs[1][1]) # trace capture + replay
	assert tt.trace_active
	pccs = {n: pcc(tt(x), r) for (n, x), r in zip(xs, refs)}
	print("[bench] PCC (trace replay, fresh input):", {k: round(v, 6) for k, v in pccs.items()})

	trace_id, dec_buf = tt._trace
	dev = tt.device
	# pure device time: back-to-back replays
	for _ in range(5):
	ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
	ttnn.synchronize_device(dev)
	t0 = time.perf_counter()
	for _ in range(a.iters):
	ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
	ttnn.synchronize_device(dev)
	dev_ms = (time.perf_counter() - t0) / a.iters * 1e3
	# 10 chunks of back-to-back replays -> median / min of the chunk means (noise on the shared host)
	chunks = []
	for _ in range(10):
	t0 = time.perf_counter()
	for _ in range(20):
	ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
	ttnn.synchronize_device(dev)
	chunks.append((time.perf_counter() - t0) / 20 * 1e3)
	dev_med, dev_min = statistics.median(chunks), min(chunks)
	print(f"[bench] device_trace_ms={dev_ms:.3f} chunk median={dev_med:.3f} min={dev_min:.3f}")
	if not tt.fused_opts.device_patch_embed: # stage loop below assumes the fused im2col input
	print("RESULT " + json.dumps({"dispatch": a.dispatch, "grid": f"{g.x}x{g.y}", "device_trace_ms": dev_ms, "pcc": pccs}))
	tt.close()
	return

	# served path, stage by stage, fresh input each request
	u8 = tt.fused_opts.input == "u8" # served input = the uint8 crop (server: preprocess_pil(raw=True))
	hp = tt._head_params
	io = tt._io
	if io is not None:
	# round 13 fast host I/O (TtHamer._forward_fast, stage by stage): prep = centring into the persistent host
	# buffer, h2d = copy_host_to_device_tensor, exec = replay + readback into the persistent host tensor + the
	# wait (spin on a CQ event, or a blocking read with HAMER_SPIN_WAIT=0), finalize = HostFinalizer
	st = {k: [] for k in ("prep", "h2d", "exec", "finalize", "total")}
	for i in range(a.iters):
	x = xs_raw[i % len(xs)] if u8 else xs[i % len(xs)][1]
	c0 = time.perf_counter()
	tt._write_input(io, x) # centring (round 14: + the tiled im2col with TT_FUSED_TILE_IN)
	c1 = time.perf_counter()
	ttnn.copy_host_to_device_tensor(io["host_in"], tt._tokens_dev)
	c2 = time.perf_counter()
	ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
	if tt.spin_wait:
	ttnn.copy_device_to_host_tensor(dec_buf, io["host_out"], blocking=False)
	ttnn.event_synchronize(ttnn.record_event(dev, 0))
	else:
	ttnn.copy_device_to_host_tensor(dec_buf, io["host_out"], blocking=True)
	c3 = time.perf_counter()
	out = io["fin"](io["out_np"])
	c4 = time.perf_counter()
	for k, (s_, e_) in zip(st, [(c0, c1), (c1, c2), (c2, c3), (c3, c4), (c0, c4)]):
	st[k].append((e_ - s_) * 1e3)
	else:
	st = {k: [] for k in ("im2col", "h2d", "exec", "d2h", "finalize", "total")}
	for i in range(a.iters):
	x = xs_raw[i % len(xs)] if u8 else xs[i % len(xs)][1]
	c0 = time.perf_counter()
	host, hdt, lay = tt._host_input(x)
	c1 = time.perf_counter()
	h = ttnn.from_torch(host, device=None, layout=lay, dtype=hdt, # round 15: the input layout (TILE with TT_FUSED_TILE_IN)
	mesh_mapper=ttnn.ReplicateTensorToMesh(dev))
	ttnn.copy_host_to_device_tensor(h, tt._tokens_dev)
	c2 = time.perf_counter()
	ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=True)
	c3 = time.perf_counter()
	dec_h = tt._dec_host(ttnn.to_torch(dec_buf), 1)
	c4 = time.perf_counter()
	out = mh.host_finalize(dec_h, hp["dec_split"], tt._init_pose, tt._init_betas, tt._init_cam, init=tt._fin_init)
	c5 = time.perf_counter()
	for k, (s, e) in zip(st, [(c0, c1), (c1, c2), (c2, c3), (c3, c4), (c4, c5), (c0, c5)]):
	st[k].append((e - s) * 1e3)
	med = {k: statistics.median(v) for k, v in st.items()}
	# last output sanity vs reference
	last = (a.iters - 1) % len(xs)
	med_pcc_last = pcc(out, refs[last])
	# full TtHamer.__call__ (what the server calls), fresh input
	call = []
	with torch.no_grad():
	for i in range(a.iters):
	x = xs_raw[i % len(xs)] if u8 else xs[i % len(xs)][1]
	c0 = time.perf_counter(); tt(x); call.append((time.perf_counter() - c0) * 1e3)
	res = {
	"dispatch": a.dispatch, "grid": f"{g.x}x{g.y}", "device_trace_ms": dev_ms,
	"device_trace_chunk_median_ms": dev_med, "device_trace_chunk_min_ms": dev_min,
	"stages_ms": med, "ttHamer_call_ms": statistics.median(call), "ttHamer_call_min_ms": min(call),
	"pcc": pccs, "pcc_bench_last": med_pcc_last,
	}
	print(f"[bench] device_trace_ms={dev_ms:.3f} (med {dev_med:.3f} min {dev_min:.3f}) served-call={res['ttHamer_call_ms']:.3f} ms (min {res['ttHamer_call_min_ms']:.3f})")
	print("[bench] stages (median ms): " + " ".join(f"{k}={v:.3f}" for k, v in med.items()))
	print("RESULT " + json.dumps(res))
	if a.json:
	Path(a.json).write_text(json.dumps(res, indent=2))
	tt.close()


	if __name__ == "__main__":
	main()