code/scripts/bench.py

#!/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.<out>.copy``     the readback split: the device-to-host copy into a preallocated host tensor ...
``d2h.<out>.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/<stem>.npz`` and
``samples/<stem>.<variant>.reference.json`` (the container-smoke gate), other inputs with ``--golden PATH`` (an
``ort_fp32.npz`` of ``research/sceneseg/public_data/golden``: ``cls_<variant>``; ``-`` 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_<variant> 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()