File size: 9,109 Bytes
ac1609a bbad5a2 ac1609a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device profile of StreamPETR's traced ``frame`` replay between signposts (OPT_BASELINE.md).
Tracy + device profiler (absolute ``-o`` under ``$ROOT/generated/profiler``, never inside the bundle):
bin/devrun -t 3600 -- python -m tracy -r -p -v --no-web-server --op-support-count 6000 \\
-o $ROOT/generated/profiler/streampetr_baseline code/scripts/profile_ops.py --rounds 2 --eager
python code/scripts/profile_summary.py $ROOT/generated/profiler/streampetr_baseline --json ... --md ...
tt-perf-report <reports/<ts>/ops_perf_results_<ts>.csv> --start-signpost frame --end-signpost frame_end --arch p150b
ttnn-visualizer memory + graph report of one eager frame (slow; never used for timing):
bin/devrun -t 3600 -- python code/scripts/profile_ops.py --graph-report $ROOT/generated/ttnn_reports/streampetr_baseline
ttnn-visualizer --profiler-path $ROOT/generated/ttnn_reports/streampetr_baseline \\
--performance-path $ROOT/generated/profiler/streampetr_baseline/reports/<ts>
The model is loaded like the API (the ``frame`` trace captured), four frames of the PandaSet sample (or synthetic
frames when it is absent) bring the stream's temporal memory to the steady state, the profiler buffer is flushed,
then the ``frame`` trace is replayed once between the signposts ``frame`` .. ``frame_end`` (inside ``trace`` ..
``trace_end``) and ``ttnn.ReadDeviceProfiler`` reads the buffer. ``--rounds 2`` repeats the window (``frame2``) to
check that the profile is stable; ``--eager`` adds one eager (untraced, program cache hot) frame between
``eager_frame`` / ``eager_frame_end`` with module-stage signposts (``bb.stem`` ... ``head.update``; the op sequence
is matched op for op against the traced one by ``profile_summary.py``).
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
HERE = Path(__file__).resolve()
sys.path.insert(0, str(HERE.parents[1]))
sys.path.insert(0, str(HERE.parent))
class _Marked:
"""Transparent wrapper of a callable: emits a Tracy signpost, then calls it (host-side only)."""
def __init__(self, obj, name: str, signpost):
self._obj, self._name, self._signpost = obj, name, signpost
def __call__(self, *a, **k):
self._signpost(self._name)
return self._obj(*a, **k)
def __getattr__(self, item):
return getattr(self._obj, item)
def mark_stages(tt, signpost):
"""Signposts at the module boundaries of the ``frame`` function (one eager run). Returns an undo callable.
Labels: bb.stem, bb.stage2 .. bb.stage5 (each from its max-pool on), bb.neck; head.tokens, head.cross_kv,
head.temporal, head.queries, head.dec<L>.self / .cross / .ffn (L = 0..5; each label starts at that block's first
projection, so the preceding residual add belongs to the previous label), head.branches, head.select,
head.update."""
undo = []
def set_attr(obj, name, value):
had = name in obj.__dict__
old = obj.__dict__.get(name)
setattr(obj, name, value)
undo.append(lambda: setattr(obj, name, old) if had else obj.__dict__.pop(name, None))
bb, hd = tt.backbone, tt.head
stem0 = bb.stem[0]
bb.stem[0] = _Marked(stem0, "bb.stem", signpost)
undo.append(lambda: bb.stem.__setitem__(0, stem0))
b20 = bb.stages[2][0]
bb.stages[2][0] = _Marked(b20, "bb.stage2", signpost)
undo.append(lambda: bb.stages[2].__setitem__(0, b20))
pool, count = bb.pool, [2]
def marked_pool(x):
count[0] += 1
signpost(f"bb.stage{count[0]}")
return pool(x)
set_attr(bb, "pool", marked_pool)
set_attr(bb, "lateral0", _Marked(bb.lateral0, "bb.neck", signpost))
for name in ("tokens", "cross_kv", "temporal", "queries", "branches", "select", "update"):
set_attr(hd, name, _Marked(getattr(hd, name), f"head.{name}", signpost))
for li, L in enumerate(hd.layers):
for key, lab in (("self.q", "self"), ("cross.q", "cross"), ("ffn.0", "ffn")):
old = L[key]
L[key] = _Marked(old, f"head.dec{li}.{lab}", signpost)
undo.append(lambda L=L, key=key, old=old: L.__setitem__(key, old))
def restore():
for fn in reversed(undo):
fn()
return restore
def steady_requests(model):
"""Five frames of one stream: the PandaSet sample f40-43 (+ f40 of the next cycle) when present, else synthetic
frames; the first four run through the API, the fifth gives the eager / graph-report inputs."""
import bench
d = bench.pandaset_dir()
paths = [d / f"pandaset_019_f{n}.json" for n in (40, 41, 42, 43)]
frames = bench.manifest_frames(paths) if all(p.is_file() for p in paths) else bench.synthetic_frames("pandaset_019")
seq = bench.Sequencer(frames, "profile")
return [seq(i) for i in range(5)]
def frame_inputs(model, req):
"""Host values of one frame for ``run_eager`` (the API's ``_forward`` without the replay)."""
from tt_streampetr.tt.params import age_input
prepared = model._prepare(**req)
plan = prepared["plan"]
params = model.tt.geometry_params(plan.M)
params.update(prepared["calib"]["params"])
inputs = {"img": model._img_input(prepared), "ages": age_input(plan.pre["pre_memory_timestamp"], device_shape=True)}
return inputs, params
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
ap.add_argument("--rounds", type=int, default=1)
ap.add_argument("--eager", action="store_true")
ap.add_argument("--graph-report", default=None, help="write a ttnn-visualizer memory/graph report of one eager "
"frame to this directory instead of profiling")
ap.add_argument("--dispatch", default=None, choices=[None, "eth", "worker"])
ap.add_argument("--num-cqs", type=int, default=None)
ap.add_argument("--weights-dir", default=None)
ap.add_argument("--json", default=None, help="write the run's metadata (config, trace description)")
a = ap.parse_args()
import ttnn
from tt_streampetr import StreamPETR
from tt_streampetr.ttaw.profiling import read_device_profiler, signpost
t0 = time.perf_counter()
model = StreamPETR.from_pretrained(dispatch=a.dispatch, num_command_queues=a.num_cqs, weights_dir=a.weights_dir)
meta = {"config": model.device_info, "knobs": model.knobs.as_dict(), "load_s": round(time.perf_counter() - t0, 1)}
print("config:", json.dumps(meta["config"]), flush=True)
dev, r = model.device, model.runner
sync = lambda: ttnn.synchronize_device(dev) # noqa: E731
try:
reqs = steady_requests(model)
for req in reqs[:4]: # steady state: a cold start, then three carried frames
model(**req)
sync()
inputs, params = frame_inputs(model, reqs[4])
if a.graph_report:
out_dir = Path(a.graph_report)
out_dir.mkdir(parents=True, exist_ok=True)
r.run_eager("frame", inputs, params) # program cache hot: nothing compiles inside the capture
sync()
t1 = time.perf_counter()
with ttnn.graph.full_graph_capture(str(out_dir / "graph_capture.json")):
r.run_eager("frame", inputs, params)
sync()
meta["graph_capture_s"] = round(time.perf_counter() - t1, 1)
from ttnn import graph_report
db = graph_report.import_report(out_dir, out_dir)
ttnn.save_config_to_json_file(str(out_dir / "config.json"))
meta["graph_report"] = {"dir": str(out_dir), "db": str(db)}
raw = out_dir / "graph_capture.json"
if raw.is_file():
meta["graph_report"]["raw_json_bytes"] = raw.stat().st_size
raw.unlink()
print("graph report:", json.dumps(meta["graph_report"]), flush=True)
else:
read_device_profiler(dev) # drop the load / warm-up data
for k in range(1, a.rounds + 1):
tag = "" if k == 1 else str(k)
signpost(f"trace{tag}")
signpost(f"frame{tag}")
r.replay("frame")
sync()
signpost(f"frame{tag}_end")
read_device_profiler(dev)
signpost(f"trace{tag}_end")
if a.eager:
restore = mark_stages(model.tt, signpost)
try:
signpost("eager_frame")
r.run_eager("frame", inputs, params)
sync()
signpost("eager_frame_end")
read_device_profiler(dev)
finally:
restore()
meta["trace"] = r.describe()
finally:
model.close()
if a.json:
Path(a.json).parent.mkdir(parents=True, exist_ok=True)
Path(a.json).write_text(json.dumps(meta, indent=1, default=str) + "\n")
return 0
if __name__ == "__main__":
sys.exit(main())
|