File size: 9,962 Bytes
5662e7f 1e8dada 5662e7f 1e8dada 5662e7f | 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 201 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device profile of ptv3-p150: per trace one eager forward (module / layer signposts) and one traced replay.
ROOT=/home/ubuntu/experiments/tt-models
$ROOT/bin/devrun -t 3600 -- python -m tracy -r -p -v --no-web-server --op-support-count 6000 \\
-o $ROOT/generated/profiler/ptv3_baseline code/scripts/profile_ops.py --input <cloud>
tt-perf-report <ops_perf_results_*.csv> --start-signpost trace_a --end-signpost trace_a_end --tracing-mode
python code/scripts/profile_summary.py <the same csv> profile_summary.json --md profile_summary.md
Sections (signposts): ``eager_a`` .. ``eager_a_end`` (``TraceRunner.run_eager`` of ``encseg_<b>``: the same graph
and programs as trace A, nothing compiled), ``trace_a`` .. ``trace_a_end`` (``--replays`` replays of the captured
trace A), then the same for trace B (``eager_b`` / ``trace_b``, variant ``seg_det``). The device profiler buffer is
flushed (``ttnn.ReadDeviceProfiler``) before and after each section, after every replay and, inside the eager
sections, before every module signpost (``--no-flush-modules`` turns that off).
Module signposts (``m:``): ``embed``, ``enc<s>.pool``, ``enc<s>.block<b>``, ``dec<s>`` (up / skip linears, unpool
gather, sum), ``dec<s>.block0``, ``seg.head`` (classifier + softmax), ``seg.summary``; ``det.fusion``,
``det.point_proj_bev`` (point_proj + the 9 BEV candidate gathers + max), ``det.bev_refine``, ``det.stem``,
``det.block<i>``, ``det.shared``, ``det.heatmap`` (hidden + out conv + sigmoid), ``det.local_max``,
``query.<part>``. Layer signposts (``c:``) inside each PTv3 block: ``cpe`` (CPE table + gather-GEMM + LN + add),
``attn.in`` (LN, two-term sequence table, gather into patch order), ``attn.qkv`` (fp32 q / k / v projections),
``attn.core`` (head split, ``q k^T``, softmax, ``p v``, head merge), ``attn.out`` (proj + inverse gather + add),
``mlp`` (LN, fc1 + GELU, fc2, add). The replay has the same op sequence as the eager forward (checked by
``profile_summary.py``), so every replayed op inherits its module / layer. The output folder must be absolute and
outside the bundle (PLAN.md 5.1). (The baseline profile of OPT_BASELINE.md ran before ``det.block<i>`` got its
``c:conv2+residual`` signpost, so its ``conv1`` layer there is the whole block.)
"""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
from typing import Callable, List, Sequence
import numpy as np
from tt_ptv3 import PTv3
from tt_ptv3.ttaw.profiling import read_device_profiler, signpost, signposted
CODE = Path(__file__).resolve().parents[1]
SAMPLES = CODE / "tt_ptv3" / "samples"
class _Signposted:
"""A callable stand-in that emits ``pre`` signposts (after an optional flush), calls the module, then emits
``post`` signposts; attributes forward to the module."""
def __init__(self, inner, pre: Sequence[str] = (), flush: Callable[[], None] = None, post: Sequence[str] = ()):
self.inner, self._pre, self._post, self._flush = inner, list(pre), list(post), flush
def __call__(self, *args, **kwargs):
if self._flush is not None and self._pre:
self._flush()
for label in self._pre:
signpost(label)
out = self.inner(*args, **kwargs)
for label in self._post:
signpost(label)
return out
def __getattr__(self, name):
return getattr(self.inner, name)
def module_signposts(graph, flush: Callable[[], None] = None) -> Callable[[], None]:
"""Wrap the modules of a ``tt.model.TtPTv3`` with signposting stand-ins; returns the restore function."""
undo: List[Callable[[], None]] = []
def swap(obj, attr, pre=(), fl=flush, post=()):
inner = getattr(obj, attr)
had = attr in getattr(obj, "__dict__", {})
setattr(obj, attr, _Signposted(inner, pre, fl, post))
undo.append((lambda: setattr(obj, attr, inner)) if had else (lambda: delattr(obj, attr)))
def swap_item(seq, i, pre=(), fl=flush):
inner = seq[i]
seq[i] = _Signposted(inner, pre, fl)
undo.append(lambda: seq.__setitem__(i, inner))
def block(blk, label):
swap(blk, "attention", ["c:attn.in"], None)
swap(blk, "q", ["c:attn.qkv"], None)
swap(blk, "v", [], None, post=["c:attn.core"])
swap(blk, "proj", ["c:attn.out"], None)
swap(blk, "norm2", ["c:mlp"], None)
swap(graph, "embed", ["m:embed"])
for s in sorted(graph.pools):
swap_item(graph.pools, s, [f"m:enc{s}.pool"])
for s, bl in graph.blocks.items():
for b in range(len(bl)):
block(bl[b], f"enc{s}.block{b}")
swap_item(bl, b, [f"m:enc{s}.block{b}", "c:cpe"])
for s in sorted(graph.dec_up):
swap_item(graph.dec_up, s, [f"m:dec{s}", "c:up"])
swap_item(graph.dec_skip, s, ["c:skip+unpool+sum"], None)
for s in sorted(graph.dec_block):
block(graph.dec_block[s], f"dec{s}.block0")
swap_item(graph.dec_block, s, [f"m:dec{s}.block0", "c:cpe"])
swap(graph, "classifier", ["m:seg.head"])
swap(graph, "seg_summary", ["m:seg.summary"])
if graph.det is not None:
swap(graph, "fusion_skip", ["m:det.fusion"])
swap(graph, "point_proj", ["m:det.point_proj_bev"])
swap(graph, "bev_refine", ["m:det.bev_refine"])
swap(graph, "stem", ["m:det.stem"])
for i, (c1, c2) in enumerate(list(graph.det_blocks)):
w1 = _Signposted(c1, [f"m:det.block{i}", "c:conv1"], flush)
if hasattr(c2, "conv_sum"): # two-term form: tt.dense.two_term_residual calls conv2.conv_sum
swap(c2, "conv_sum", ["c:conv2+residual"], None)
w2 = c2
else:
w2 = _Signposted(c2, ["c:conv2+residual"], None)
graph.det_blocks[i] = (w1, w2)
undo.append(lambda i=i, c1=c1, c2=c2: graph.det_blocks.__setitem__(i, (c1, c2)))
swap(graph, "shared", ["m:det.shared"])
swap(graph, "heat_hidden", ["m:det.heatmap"])
swap(graph, "local_max", ["m:det.local_max"])
qh = graph.query_head
swap(qh.select, "indices", ["m:query.topk"])
swap(qh.select, "decode", ["m:query.decode"])
for attr, label in (("init_queries", "m:query.init"), ("keys", "m:query.keys"),
("decoder", "m:query.decoder"), ("predict", "m:query.predict"),
("query_heat", "m:query.heat"), ("_indices_tile", "m:query.pack")):
swap(qh, attr, [label])
def restore() -> None:
for fn in reversed(undo):
fn()
return restore
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
ap.add_argument("--input", default=None, help="cloud (default: the shipped sample)")
ap.add_argument("--dispatch", default=None, choices=["eth", "worker"])
ap.add_argument("--cqs", type=int, default=None, choices=[1, 2])
ap.add_argument("--variant", default=None, choices=["seg_det", "seg"])
ap.add_argument("--replays", type=int, default=1, help="replays inside each trace section")
ap.add_argument("--weights-dir", default=None)
ap.add_argument("--no-eager", dest="eager", action="store_false")
ap.add_argument("--no-flush-modules", dest="flush_modules", action="store_false")
a = ap.parse_args()
import ttnn
from tt_ptv3 import io as tio
from tt_ptv3.host import preprocess as pre
from tt_ptv3.tt import inputs as tin
src = a.input
if src is None:
src = next(str(SAMPLES / n) for n in ("test_pcd.npz", "pandaset_019_f40.npz") if (SAMPLES / n).is_file())
t0 = time.perf_counter()
with PTv3.from_pretrained(dispatch=a.dispatch, num_command_queues=a.cqs, variant=a.variant,
weights_dir=a.weights_dir) as model:
tr, dev = model.runner, model.device
print("loaded in %.1f s:" % (time.perf_counter() - t0), json.dumps(model.device_info, default=str), flush=True)
out = model(src) # one served frame: upload + replays + readbacks
print(json.dumps({"input": src, "detections": out.to_dict().get("num_detections"),
"frame": out.meta.get("frame")}, default=str), flush=True)
frame = pre.prepare_frame(np.asarray(tio.load_points(src).points, np.float32))
b, caps = model.bucket_for(frame.counts) # the bucket the served frame ran in
print(json.dumps({"bucket": b, "caps": list(caps)}), flush=True)
tables = {f"{b}.{k}": v for k, v in tin.frame_inputs(frame, caps, det=model.with_det).items()}
ttnn.synchronize_device(dev)
read_device_profiler(dev) # warm-up / capture / first frame out of the buffer
def flush():
ttnn.synchronize_device(dev)
read_device_profiler(dev)
variants = [("a", f"encseg_{b}")] + ([("b", f"det_{b}")] if model.with_det else [])
if a.eager:
restore = module_signposts(model.graph, flush=flush if a.flush_modules else None)
try:
for tag, v in variants:
with signposted(f"eager_{tag}"):
tr.run_eager(v, inputs=tables if tag == "a" else None)
ttnn.synchronize_device(dev)
read_device_profiler(dev)
finally:
restore()
tr.upload(tables)
ttnn.synchronize_device(dev)
read_device_profiler(dev)
for tag, v in variants:
with signposted(f"trace_{tag}"):
for _ in range(a.replays):
tr.replay(v)
ttnn.synchronize_device(dev)
read_device_profiler(dev)
read_device_profiler(dev)
print("trace:", json.dumps(tr.describe(), default=str), flush=True)
if __name__ == "__main__":
main()
|