| |
| |
| """Phase-split benchmark of the TT_FUSED serving path (baseline profiling helper). |
| |
| Run from code/ after `source tools/chipenv.sh <chip> <venv>`; PYTHONPATH must start with code/. |
| SP_DISPATCH=auto|eth|worker (default auto = ETH, 1 CQ) SP_N_ITER=200 (compute grid is always 12x10) python bench_fused.py |
| Measures, each as a median over N iterations with a device sync per iteration (or a pipelined |
| mean where noted): |
| trace_only execute_trace + synchronize (input already resident) -> device forward |
| trace_pipelined N back-to-back execute_trace, one sync -> device throughput |
| h2d copy_host_to_device_tensor of the prepared 614 KB input + sync |
| d2h_desc to_torch of the resident descriptor output (RM bf16 [1,1,4800,256]) |
| d2h_nms to_torch of the resident NMS map (RM bf16 [1,1,480,640]) |
| host_convert reshape/permute/float of the readbacks |
| forward run_fused() end to end (prepare host input + H2D + trace + D2H + convert) |
| post postprocess_from_nms_map (threshold/border/top-k/grid_sample) |
| SP_EAGER_PROFILE=1: run only the eager fused graph a few times (for the device profiler; no trace). |
| """ |
| import os |
| import statistics |
| import time |
|
|
| import torch |
| import ttnn |
|
|
| from models.reference.superpoint_reference import get_natural_input, load_reference_model |
| from models.tt import postprocess as _post |
| from models.tt.superpoint_ttnn import TtSuperPoint |
|
|
|
|
| def _open(): |
| from models.tt.device_open import open_ttnn_device, resolve_dispatch |
|
|
| mode = resolve_dispatch(os.environ.get("SP_DISPATCH", "auto")) |
| kw = dict(l1_small_size=32 * 1024, trace_region_size=32 << 20) |
| return open_ttnn_device(0, dispatch=mode, **kw), mode |
|
|
|
|
| def _med(f, n): |
| ts = [] |
| for _ in range(n): |
| t0 = time.perf_counter() |
| f() |
| ts.append((time.perf_counter() - t0) * 1e3) |
| return statistics.median(ts), min(ts) |
|
|
|
|
| def main(): |
| n = int(os.environ.get("SP_N_ITER", "200")) |
| dev, mode = _open() |
| g = dev.compute_with_storage_grid_size() |
| try: |
| tm = load_reference_model() |
| px = get_natural_input(batch_size=1) |
| sp = TtSuperPoint(tm, dev, fused=True) |
| tt_in = sp.allocate_input(1) |
| if os.environ.get("SP_EAGER_PROFILE") == "1": |
| sp.load_input(tt_in, px) |
| for _ in range(3): |
| outs = sp.build_fused_graph(tt_in, 1) |
| ttnn.synchronize_device(dev) |
| outs.deallocate() |
| ttnn.ReadDeviceProfiler(dev) |
| print("eager profile done") |
| return |
| sp.run_fused(tt_in, px) |
| ttnn.synchronize_device(dev) |
| sp.capture_trace(tt_in, 1) |
| host_in = sp.prepare_host_input(px) |
| sp.load_input_prepared(tt_in, host_in) |
| ttnn.synchronize_device(dev) |
| tid = sp.trace_id |
| outs = sp._trace_outputs |
| for _ in range(10): |
| ttnn.execute_trace(dev, tid, cq_id=0, blocking=False) |
| ttnn.synchronize_device(dev) |
|
|
| def trace_once(): |
| ttnn.execute_trace(dev, tid, cq_id=0, blocking=False) |
| ttnn.synchronize_device(dev) |
|
|
| r = {} |
| r["trace_only"] = _med(trace_once, n) |
| t0 = time.perf_counter() |
| for _ in range(n): |
| ttnn.execute_trace(dev, tid, cq_id=0, blocking=False) |
| ttnn.synchronize_device(dev) |
| r["trace_pipelined"] = ((time.perf_counter() - t0) * 1e3 / n,) * 2 |
|
|
| def h2d(): |
| ttnn.copy_host_to_device_tensor(host_in, tt_in, cq_id=0) |
| ttnn.synchronize_device(dev) |
|
|
| r["h2d"] = _med(h2d, n) |
| r["prepare_host_input"] = _med(lambda: sp.prepare_host_input(px), n) |
| r["d2h_desc"] = _med(lambda: ttnn.to_torch(outs.d_out), n) |
| r["d2h_nms"] = _med(lambda: ttnn.to_torch(outs.nms_map), n) |
| r["d2h_scores_fallback"] = _med(lambda: ttnn.to_torch(outs.s_sm), max(10, n // 10)) |
| r["read_fused_total"] = _med(lambda: sp._read_fused(outs, 1, 480, 640, sp.nms_radius_traced), n) |
| r["forward"] = _med(lambda: sp.run_fused(tt_in, px), n) |
| res = sp.run_fused(tt_in, px) |
| r["post"] = _med( |
| lambda: _post.postprocess_from_nms_map( |
| res.nms_map, res.descriptors_nchw, keypoint_threshold=0.005, max_keypoints=1024, |
| border_removal_distance=4, with_descriptors=True, |
| ), |
| n, |
| ) |
| |
| |
| kw = dict(keypoint_threshold=0.005, max_keypoints=1024, border_removal_distance=4) |
|
|
| def e2e_legacy_api(): |
| r_ = sp.run_fused(tt_in, px) |
| return _post.postprocess_from_nms_map(r_.nms_map, r_.descriptors_nchw, with_descriptors=True, **kw)[0] |
|
|
| r["e2e_run_fused+post"] = _med(e2e_legacy_api, n) |
| if hasattr(sp, "run_fused_keypoints"): |
| r["e2e_keypoints"] = _med(lambda: sp.run_fused_keypoints(tt_in, sp.prepare_host_input(px), **kw), n) |
| a = e2e_legacy_api() |
| b = sp.run_fused_keypoints(tt_in, sp.prepare_host_input(px), **kw) |
| same = torch.equal(a[0], b[0]) and torch.equal(a[1], b[1]) |
| print(f"e2e_keypoints vs run_fused+post: kp/scores identical={same} n_kp={a[0].shape[0]} " |
| f"desc max|diff|={float((a[2] - b[2]).abs().max()):.2e}") |
| if hasattr(sp, "run_fused_keypoints_kpc") and sp._gather_tid is not None: |
| r["e2e_kpc"] = _med(lambda: sp.run_fused_keypoints_kpc(tt_in, sp.prepare_host_input(px), **kw), n) |
| a = e2e_legacy_api() |
| c = sp.run_fused_keypoints_kpc(tt_in, sp.prepare_host_input(px), **kw) |
| b = sp.run_fused_keypoints(tt_in, sp.prepare_host_input(px), **kw) |
| print(f"e2e_kpc vs run_fused+post: kp/scores identical={torch.equal(a[0], c[0]) and torch.equal(a[1], c[1])} " |
| f"n_kp={c[0].shape[0]} desc max|diff|={float((a[2] - c[2]).abs().max()):.2e}; " |
| f"vs run_fused_keypoints: all identical={all(torch.equal(x, y) for x, y in zip(b, c))}") |
| if hasattr(sp, "prepare_host_input_u8") and sp._gather_tid is not None: |
| |
| r8 = torch.round(px[0, 0] * 255.0).to(torch.uint8).numpy() |
| r["e2e_kpc_u8"] = _med(lambda: sp.run_fused_keypoints_kpc(tt_in, sp.prepare_host_input_u8(r8), **kw), n) |
| print(f"mode={mode} grid={g.x}x{g.y} n_iter={n}") |
| for k, (med, mn) in r.items(): |
| print(f"{k:22s} median={med:8.3f} ms min={mn:8.3f} ms") |
| sp.release() |
| finally: |
| ttnn.close_device(dev) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
