| |
| """Device bring-up of the planner graph (development tool; needs the p150, run through ``bin/devrun``). |
| |
| bin/devrun -t 1800 -- python code/scripts/bringup_device.py [--scenes kashiwanoha_dense straight_road] |
| [--no-capture] [--precision "dec.*=HiFi2+fp32"] [--json logs/diffusion-planner/bringup.json] |
| |
| Protocol (PLAN.md 4.4): build ``TtDiffusionPlanner(debug=True)``; per scene run every variant EAGERLY twice |
| (``encoder_taps``, ``decode_once`` at evaluations 0 / 5 / 10, ``plan``), check the two eager runs are bit-identical and |
| compare them with the research goldens (``research/diffusion-planner/goldens/<scene>.npz``, the fp32 CPU reference: |
| PCC on valid rows, max abs; final x0 on valid agents; logits); then capture all variants (strict, no program |
| compiled after the capture) and check replay == eager bit for bit on the same inputs. Prints a table and writes a |
| JSON report. This is a diagnostic: the frozen gates live in ``tests/test_pcc_device.py`` / ``test_e2e_device.py``. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| HERE = Path(__file__).resolve() |
| sys.path.insert(0, str(HERE.parents[1])) |
|
|
| from tt_diffusion_planner.host import pipeline as hp |
| from tt_diffusion_planner.reference import config as C |
| from tt_diffusion_planner.ttaw.metrics import error_stats |
|
|
| GOLDENS = HERE.parents[4] / "research" / "diffusion-planner" / "goldens" |
|
|
|
|
| def stats(dev, ref): |
| s = error_stats(np.asarray(dev, np.float64), np.asarray(ref, np.float64)) |
| return {"pcc": round(s["pcc"], 7), "max_abs": float(f"{s['max_abs']:.4g}"), "rel_l2": float(f"{s['rel_l2']:.4g}")} |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) |
| ap.add_argument("--scenes", nargs="+", default=["kashiwanoha_dense", "straight_road"]) |
| ap.add_argument("--no-capture", action="store_true") |
| ap.add_argument("--precision", default=None) |
| ap.add_argument("--ln-fp32", default=None, help="module globs (comma-separated) for the fp32 LayerNorm") |
| ap.add_argument("--hidden-fp32", default=None, help="module globs (comma-separated) for fp32 hidden activations") |
| ap.add_argument("--split", default=None, help="module globs (comma-separated) for split (hi/lo) fp32 matmuls") |
| ap.add_argument("--attn-fp32-acc", default=None, help="module globs for SDPA with fp32 accumulation") |
| ap.add_argument("--attn-matmul", default=None, help="module globs for the fp32 matmul attention") |
| ap.add_argument("--plan-only", action="store_true", help="skip the encoder taps and decoder evaluations") |
| ap.add_argument("--evals", nargs="+", type=int, default=[0, 5, 10]) |
| ap.add_argument("--json", default=None) |
| args = ap.parse_args() |
|
|
| import torch |
|
|
| torch.set_num_threads(4) |
| from tt_diffusion_planner.device import close_device, open_device |
| from tt_diffusion_planner.reference.weights import find_weights_dir, load_weights |
| from tt_diffusion_planner.tt.model import TtDiffusionPlanner |
|
|
| weights = load_weights(find_weights_dir()) |
| dev = open_device(allow_fallback=False) |
| g = dev.compute_with_storage_grid_size() |
| report = {"grid": f"{g.x}x{g.y}", "precision": args.precision, "scenes": {}} |
| t0 = time.perf_counter() |
| from tt_diffusion_planner.tt.config import globs |
|
|
| opts = {k: globs(v) for k, v in (("ln_fp32", args.ln_fp32), ("hidden_fp32", args.hidden_fp32), |
| ("split", args.split), ("attn_fp32_acc", args.attn_fp32_acc), |
| ("attn_matmul", args.attn_matmul)) if v is not None} |
| tt = TtDiffusionPlanner(dev, weights, debug=True, precision=args.precision, **opts) |
| report["options"] = tt.build.options() |
| report["build_s"] = round(time.perf_counter() - t0, 2) |
| print(f"build {report['build_s']} s, uploaded {tt.build.uploaded_bytes / 2**20:.1f} MiB", flush=True) |
| runner = tt.runner |
| from tt_diffusion_planner.tt import inputs as I |
|
|
| eager_cache = {} |
| try: |
| for scene in args.scenes: |
| with np.load(GOLDENS / f"{scene}.npz", allow_pickle=False) as z: |
| gold = {k: z[k] for k in z.files if k != "__meta__"} |
| raw = {k: gold[f"in.{k}"] for k in C.INPUT_NAMES} |
| prep = hp.prepare(raw, weights.normalization.observation) |
| res = {} |
| if args.plan_only: |
| args.evals = [] |
| |
| t1 = time.perf_counter() |
| taps = tt.encoder_taps(prep, eager=True) |
| taps2 = taps if args.plan_only else tt.encoder_taps(prep, eager=True) |
| res["encoder_eager_s"] = round(time.perf_counter() - t1, 2) |
| res["encoder_deterministic"] = all(np.array_equal(taps[k], taps2[k]) for k in taps) |
| for name, _ in C.TOKEN_LAYOUT: |
| rows = np.flatnonzero(gold[f"host.valid.{name}"]) |
| if rows.size: |
| res[f"enc.{name}"] = stats(taps[f"enc.{name}"][rows], gold[f"enc.{name}"][rows]) |
| for c in ("ego", "neighbor", "lane", "route", "polygon", "line_string"): |
| rows = gold[f"enc.{c}.pre.rows"] |
| if rows.size: |
| for part in ("pre", "mixer"): |
| res[f"enc.{c}.{part}"] = stats(taps[f"enc.{c}.{part}"][rows], gold[f"enc.{c}.{part}"]) |
| tok = np.flatnonzero(gold["host.token_valid"]) |
| for name in ["enc.tokens"] + [f"enc.fusion.{i}" for i in range(6)] + ["enc.encoding"]: |
| res[name] = stats(taps[name][tok], gold[name][tok]) |
| |
| arows = gold["dec.rows"] |
| for k in args.evals: |
| outs = [tt_decode(tt, runner, prep, gold, k) for _ in range(2)] |
| res[f"dec.eval{k}.deterministic"] = bool(np.array_equal(outs[0], outs[1])) |
| res[f"dec.eval{k}"] = stats(outs[0][arows][:, 1:], gold["dec.out"][k][:, 1:]) |
| |
| t1 = time.perf_counter() |
| p1 = runner.run_eager("plan", inputs=I.plan_inputs(prep)) |
| p2 = runner.run_eager("plan", inputs=I.plan_inputs(prep)) |
| res["plan_eager_s"] = round(time.perf_counter() - t1, 2) |
| res["plan_deterministic"] = all(np.array_equal(p1[k], p2[k]) for k in p1) |
| eager_cache[scene] = (prep, p1) |
| final = p1["final_x0"].reshape(352, 324)[:321].reshape(321, 81, 4) |
| res["plan.final_x0"] = stats(final[arows], gold["final_x0"][arows]) |
| res["plan.logit"] = {"dev": [round(float(v), 4) for v in p1["logit"].reshape(-1)[:5]], |
| "ref": [round(float(v), 4) for v in gold["turn.logit"]]} |
| out = hp.make_output(final, p1["logit"].reshape(-1)[:5], prep, weights.normalization, |
| {k: v[3] for k, v in hp.RUNTIME_PARAMS.items()}) |
| dpos = np.hypot(*(out.poses[:, :2] - gold["out.trajectory"][:, :2]).T) |
| res["plan.ego_max_m"], res["plan.ego_mean_m"] = round(float(dpos.max()), 4), round(float(dpos.mean()), 4) |
| res["plan.turn_equal"] = int(out.turn_indicator["command"]) == int(gold["out.turn_command"]) |
| if gold["out.predicted_agents"].shape[0]: |
| dxy = out.predicted_agents[..., :2] - gold["out.predicted_agents"][..., :2] |
| per = np.hypot(*dxy.transpose(2, 0, 1)) |
| res["plan.nb_median_max_m"] = round(float(np.median(per.max(axis=1))), 4) |
| report["scenes"][scene] = res |
| print(json.dumps({scene: res}, indent=1), flush=True) |
| if not args.no_capture: |
| t1 = time.perf_counter() |
| runner.capture() |
| report["capture_s"] = round(time.perf_counter() - t1, 2) |
| report["timings_ms"] = runner.timings_ms |
| for scene, (prep, p1) in eager_cache.items(): |
| rep = runner("plan", inputs=I.plan_inputs(prep)) |
| report["scenes"][scene]["plan_replay_equals_eager"] = all(np.array_equal(rep[k], p1[k]) for k in p1) |
| import ttnn |
|
|
| ttnn.synchronize_device(dev) |
| t2 = time.perf_counter() |
| runner.replay("plan", 5) |
| ttnn.synchronize_device(dev) |
| report["scenes"][scene]["plan_replay_ms"] = round((time.perf_counter() - t2) / 5 * 1e3, 3) |
| print(json.dumps({k: v for k, v in report.items() if k != "scenes"}, indent=1, default=str), flush=True) |
| print(json.dumps({s: {k: v for k, v in r.items() if k.startswith("plan_")} |
| for s, r in report["scenes"].items()}, indent=1), flush=True) |
| finally: |
| tt.release() |
| close_device(dev) |
| if args.json: |
| Path(args.json).parent.mkdir(parents=True, exist_ok=True) |
| Path(args.json).write_text(json.dumps(report, indent=1, default=str) + "\n") |
| return 0 |
|
|
|
|
| def tt_decode(tt, runner, prep, gold, k): |
| return tt.decode_once(prep, gold["dec.x_in"][k], float(gold["dec.t"][k]), encoding=gold["enc.encoding"], |
| eager=True) |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
