#!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 """Stage breakdown of warm forwards: the numbers OPT_BASELINE.md / OPT_REPORT.md / the card quote. bin/devrun -t 900 -- python code/scripts/bench.py --iters 100 --json logs/baseline_bench.json bin/devrun -t 900 -- python code/scripts/bench.py --input a.jpg --golden a_ort_fp32.npz --dispatch worker --cqs 2 For every ``--input`` image (default: the shipped sample), batch 1, after ``--warmup`` calls: ``--iters`` iterations of the served path cut into stages, each synchronised on its own (p50 / p99 / mean / min in ms, so the rows do not add up to ``e2e`` exactly), then ``--iters`` plain ``model(image)`` calls, then back-to-back replays: ===================== ================================================================================================ ``pre`` host pre-processing as served (``model._prepare``): decode + RGB -> BGR + the exact OpenCV INTER_LINEAR squash to 640x320 (``ttaw.image_linear.sceneseg_resize``) ``pre.decode`` decode + RGB -> BGR alone (``ttaw.io.load_image``) ``pre.squash`` the squash alone, on the decoded BGR8 image ``host_in`` copy of the uint8 squash (``INPUT_LAYOUT`` shape) into the persistent ``HostStaging`` buffer ``h2d`` upload of the squash (614 KB) + device synchronise ``trace`` one blocking replay of the whole-graph trace ``d2h`` ``TraceRunner.read`` as served: the class map [1, 320, 640] uint32 -> numpy ``d2h..copy`` the readback split: the device-to-host copy into a preallocated host tensor ... ``d2h..to_numpy`` ... and its host conversion to numpy ``post`` host post-processing (uint32 -> uint8, class 1 -> 255, INTER_NEAREST to the source size) ``e2e`` ``model(image)`` as served (its ``timing_ms`` medians are reported too) ``b2b`` back-to-back replays, one synchronise per round of ``--b2b-iters``: device time per forward ===================== ================================================================================================ AICLK / power / temperature are sampled from sysfs during the run (``ttaw.profiling.AiclkSampler``), and the host load average is recorded (the host stages move with it). ``--check`` (default) compares the served class map of each input with its stored fp32 CPU reference: the shipped sample with ``tests/goldens/.npz`` and ``samples/..reference.json`` (the container-smoke gate), other inputs with ``--golden PATH`` (an ``ort_fp32.npz`` of ``research/sceneseg/public_data/golden``: ``cls_``; ``-`` for none). Always quote the configuration line (dispatch / CQs / grid / knobs) with the numbers. """ from __future__ import annotations import argparse import json import os import platform import statistics import time from pathlib import Path from typing import Any, Dict, List, Optional import numpy as np from tt_sceneseg import SceneSeg from tt_sceneseg.reference import assets from tt_sceneseg.ttaw.profiling import AiclkSampler, StageBench, time_b2b CODE = Path(__file__).resolve().parents[1] SAMPLE = CODE / "tt_sceneseg" / "samples" / "highway_normal_1.png" GOLDENS = CODE / "tt_sceneseg" / "tests" / "goldens" def _golden_class_map(path: Path, golden: Optional[Path], variant: str) -> Optional[np.ndarray]: if golden is None: cand = GOLDENS / f"{path.stem}.npz" if cand.is_file(): return np.load(cand)[f"{variant}.class_map"] return None g = np.load(golden) for key in (f"cls_{variant}", f"{variant}.class_map"): if key in g: return g[key] raise KeyError(f"{golden}: no class map of {variant}") def check_output(model, path: Path, golden: Optional[Path]) -> Dict[str, Any]: """Served class map of ``path`` vs its stored fp32 CPU reference (when one exists).""" from tt_sceneseg.server import smoke_test as st from tt_sceneseg.ttaw.server import smoke variant = model.variant out = model(str(path)) res: Dict[str, Any] = {"src_hw": list(out.mask.shape), "fg_fraction": float((out.mask == 255).mean())} ref_cls = _golden_class_map(path, golden, variant) if ref_cls is not None: cls = model.runner(model._prepare(str(path))["prepared"].device_input())["class_map"] res["golden"] = str(golden or GOLDENS / f"{path.stem}.npz") res["argmax_agreement_320x640"] = float((cls == ref_cls).mean()) ref = path.parent / f"{path.stem}.{variant}.reference.json" if ref.is_file(): metrics, fails = smoke.compare_with_reference(out.to_dict(), smoke.load_json(ref), gates=st.REFERENCE_GATES) res["reference"] = {"file": str(ref), "metrics": metrics, "failures": fails} return res def bench_input(model, path: Path, *, iters: int, warmup: int, b2b_iters: int, b2b_rounds: int, name: str) -> Dict[str, Any]: import ttnn from tt_sceneseg.host.preprocess import prepare_bgr, rgb_to_bgr from tt_sceneseg.ttaw import io as tio from tt_sceneseg.ttaw.tensors import to_numpy sr = model.runner # tt.runner.SceneSegRunner tr = sr.runner # ttaw.trace.TraceRunner dev, var = model.device, sr.variant params = model.validate_params({}) src = str(path) bench = StageBench(name) for _ in range(warmup): model(src) outs = tr.outputs(var) # {"class_map": ...} device tensors of the trace hosts = {k: ttnn.allocate_tensor_on_host(t.spec, dev) for k, t in outs.items()} for _ in range(iters): t0 = time.perf_counter() prepared = model._prepare(src) t1 = time.perf_counter() host = sr._upload_value(prepared["prepared"].device_input()) t2 = time.perf_counter() tr.upload({"image": host}) ttnn.synchronize_device(dev) t3 = time.perf_counter() tr.replay(var) ttnn.synchronize_device(dev) t4 = time.perf_counter() res = tr.read(var) t5 = time.perf_counter() model._postprocess(sr._unpack(res), prepared, params) t6 = time.perf_counter() for stage, a, b in (("pre", t0, t1), ("host_in", t1, t2), ("h2d", t2, t3), ("trace", t3, t4), ("d2h", t4, t5), ("post", t5, t6)): bench.add(stage, (b - a) * 1e3) # the split rows (outside the served path; same data) with bench.stage("pre.decode"): bgr = rgb_to_bgr(tio.load_image(src)) with bench.stage("pre.squash"): prepare_bgr(bgr, channel_order=model.variant) for k, t in outs.items(): with bench.stage(f"d2h.{k}.copy"): ttnn.copy_device_to_host_tensor(t, hosts[k], blocking=True) with bench.stage(f"d2h.{k}.to_numpy"): to_numpy(hosts[k]) timing: List[Dict[str, float]] = [] for _ in range(iters): t0 = time.perf_counter() out = model(src) bench.add("e2e", (time.perf_counter() - t0) * 1e3) timing.append(dict(out.timing_ms)) rounds = [time_b2b(lambda: tr.replay(var), lambda: ttnn.synchronize_device(dev), n=b2b_iters) for _ in range(b2b_rounds)] for r in rounds: bench.add("b2b", r) b2b = statistics.median(rounds) return {"input": src, "stages_ms": bench.summary(), "table": bench.table(), "timing_ms_median": {k: statistics.median(t[k] for t in timing) for k in timing[0]}, "b2b_rounds_ms": rounds, "b2b_iters": b2b_iters, "device_fps_b2b": 1e3 / b2b, "e2e_fps_sync": 1e3 / bench.summary()["e2e"]["p50"]} def main() -> None: ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) ap.add_argument("--input", action="append", default=None, help="image (repeatable; default: the shipped sample)") ap.add_argument("--golden", action="append", default=None, help="golden ort_fp32.npz per --input (cls_ class map), '-' for the default lookup") ap.add_argument("--iters", type=int, default=100) ap.add_argument("--warmup", type=int, default=10) ap.add_argument("--b2b-iters", type=int, default=50) ap.add_argument("--b2b-rounds", type=int, default=5) ap.add_argument("--variant", default=None, help="channel order (default: SCENESEG_VARIANT or bgr)") ap.add_argument("--dispatch", default=None, choices=["eth", "worker"], help="default: SCENESEG_DISPATCH or eth") ap.add_argument("--cqs", type=int, default=None, choices=[1, 2], help="command queues (default: SCENESEG_NUM_CQS or 1)") ap.add_argument("--allow-fallback", action="store_true", help="accept a WORKER fallback of an ETH open") ap.add_argument("--no-check", dest="check", action="store_false") ap.add_argument("--tag", default="", help="label of this configuration in the output") ap.add_argument("--json") a = ap.parse_args() inputs = [Path(p) for p in (a.input or [str(SAMPLE)])] goldens = [None if g in (None, "-") else Path(g) for g in (a.golden or [])] + [None] * len(inputs) weights = assets.find_weights_dir() load0 = os.getloadavg() t0 = time.perf_counter() with SceneSeg.from_pretrained(variant=a.variant, dispatch=a.dispatch, num_command_queues=a.cqs, weights_dir=str(weights) if weights else None) as model: load_s = time.perf_counter() - t0 info = dict(model.device_info) want = a.dispatch or os.environ.get("SCENESEG_DISPATCH", "eth") if info.get("dispatch") != want and not a.allow_fallback: raise SystemExit(f"device opened with dispatch {info.get('dispatch')!r}, asked for {want!r}") tag = a.tag or f"{info.get('dispatch')}-{info.get('num_command_queues')}cq {info.get('grid')}" res: Dict[str, Any] = {"tag": tag, "config": info, "variant": model.variant, "knobs": model.knobs.as_dict(), "weights": model.weights, "model": model.runner.describe()["model"], "trace": model.runner.runner.describe(), "load_s": load_s, "warmup_ms": dict(model.warmup_ms), "iters": a.iters, "host": {"cpus": os.cpu_count(), "machine": platform.processor() or platform.machine(), "loadavg_start": load0}, "inputs": {}} print("config:", json.dumps({"tag": tag, "config": info, "knobs": res["knobs"]}, default=str), flush=True) with AiclkSampler(interval_s=0.05) as clk: for path, golden in zip(inputs, goldens): r = bench_input(model, path, iters=a.iters, warmup=a.warmup, b2b_iters=a.b2b_iters, b2b_rounds=a.b2b_rounds, name=f"{path.name} | {tag} | {model.variant}") if a.check: r["check"] = check_output(model, path, golden) res["inputs"][path.name] = r print(r["table"], flush=True) print(f"b2b rounds ms: {[round(x, 3) for x in r['b2b_rounds_ms']]} " f"timing_ms medians: { {k: round(v, 2) for k, v in r['timing_ms_median'].items()} }", flush=True) if a.check: print("check:", json.dumps(r["check"], default=str), flush=True) res["aiclk"] = clk.summary() res["host"]["loadavg_end"] = os.getloadavg() print(json.dumps({"tag": tag, "config": info, "aiclk": res["aiclk"], "load_s": round(load_s, 2), "loadavg": [load0, res["host"]["loadavg_end"]]}, default=str), flush=True) if a.json: Path(a.json).parent.mkdir(parents=True, exist_ok=True) Path(a.json).write_text(json.dumps(res, indent=1, default=str) + "\n") if __name__ == "__main__": main()