File size: 9,063 Bytes
7cc5047 | 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 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device profile of bevfusion-p150: one eager frame and one traced frame (``sparse_<bucket>`` + ``dense``) between
Tracy signposts.
ROOT=/home/ubuntu/experiments/tt-models
$ROOT/bin/devrun -t 3600 -- python -m tracy -r -p -v --no-web-server \\
-o $ROOT/generated/profiler/bevfusion_baseline code/scripts/profile_ops.py
tt-perf-report <ops_perf_results_*.csv> --start-signpost trace --end-signpost trace_end --tracing-mode
python code/scripts/profile_summary.py <the same csv> profile_summary.json --md profile_summary.md
Sections (signposts): ``eager`` .. ``eager_end``, one eager frame after the warm-up (``TraceRunner.run_eager`` of the
frame's ``sparse_<bucket>`` variant, then of ``dense``: the same graph and programs as the traces, nothing
compiled), and ``trace`` .. ``trace_end``, ``--replays`` traced frames (the served device path: the two replays,
the device profiler flushed after each; the uploaded inputs are the ``--input`` frame's). The device profiler
buffer (~1,000 programs) is also flushed (``ttnn.ReadDeviceProfiler``) before and after each section and, inside
the eager section, before every module (``--no-flush-modules`` turns that off).
Inside the eager section every module opens a signpost: ``m:sparse.<layer>`` for each of the 21 sparse layers (its
gather, matmul, bias / residual add and the next table), ``m:sparse.to_dense`` (the two to-dense gathers, the
channel concat and the state write), ``m:backbone.block<i>`` with a layer signpost ``c:backbone.block<i>.conv<j>``
per conv, ``m:neck``, ``m:head.shared``, ``m:head.heatmap.hidden``, ``m:head.heatmap``, ``m:head.sigmoid``,
``m:head.local_max``, ``m:query.topk``, ``m:query.decode``, ``m:query.init``, ``m:query.keys``, ``m:query.decoder``,
``m:query.predict``, ``m:query.heat``, ``m:query.pack``; so the ops of the eager frame, and through the identical op
order those of the replays, are attributed to modules (``profile_summary.py``). With
``TT_METAL_PROFILER_MID_RUN_DUMP=1`` the profiler also dumps mid-trace (needed only when one replay exceeds the
buffer). The output folder must be absolute and outside the bundle (PLAN.md 5.1).
"""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
from typing import Callable, List, Sequence
from tt_bevfusion import BEVFusion
from tt_bevfusion.ttaw.profiling import read_device_profiler, signpost, signposted
SAMPLE = Path(__file__).resolve().parents[1] / "tt_bevfusion" / "samples" / "test_pcd.npz"
class _Signposted:
"""A callable stand-in for a module that emits its signposts, then calls the module (attributes forward)."""
def __init__(self, inner, labels: Sequence[str], flush: Callable[[], None] = None):
self.inner, self._labels, self._flush = inner, list(labels), flush
def __call__(self, *args, **kwargs):
if self._flush is not None:
self._flush()
for label in self._labels:
signpost(label)
return self.inner(*args, **kwargs)
def __getattr__(self, name):
return getattr(self.inner, name)
def module_signposts(graph, layers: bool = True, flush: Callable[[], None] = None) -> Callable[[], None]:
"""Wrap the modules of a ``tt.model.TtBEVFusion`` (and the module-level functions its forward calls) with
signposting stand-ins; returns the function that restores them. ``flush`` runs before each module signpost."""
from tt_bevfusion.tt import model as tt_model
from tt_bevfusion.tt import sparse as tt_sparse
undo: List[Callable[[], None]] = []
def swap(obj, attr, labels, fl=flush):
inner = getattr(obj, attr)
setattr(obj, attr, _Signposted(inner, labels, fl))
undo.append(lambda: setattr(obj, attr, inner))
def swap_method(obj, attr, labels, fl=flush):
"""An instance attribute that shadows a bound method (removed again on restore)."""
inner = getattr(obj, attr)
setattr(obj, attr, _Signposted(inner, labels, fl))
undo.append(lambda: delattr(obj, attr))
def swap_item(seq, i, labels, fl=flush):
inner = seq[i]
seq[i] = _Signposted(inner, labels, fl)
undo.append(lambda: seq.__setitem__(i, inner))
def swap_fn(module, name, label):
inner = getattr(module, name)
def wrapped(*args, **kwargs):
if flush is not None:
flush()
signpost(label)
return inner(*args, **kwargs)
setattr(module, name, wrapped)
undo.append(lambda: setattr(module, name, inner))
# the sparse layers: the stand-in's __call__ signposts, its to_table (called right after) inherits the label
for i, layer in enumerate(graph.sparse.layers):
swap_item(graph.sparse.layers, i, [f"m:{layer.name}"])
swap_fn(tt_sparse, "to_dense", "m:sparse.to_dense")
for i, block in enumerate(graph.backbone.blocks):
for j in range(len(block)):
labels = ([f"m:backbone.block{i}"] if j == 0 else []) + ([f"c:backbone.block{i}.conv{j}"] if layers else [])
swap_item(block, j, labels, flush if j == 0 else None)
swap(graph, "neck", ["m:neck"])
swap(graph, "shared", ["m:head.shared"])
swap(graph, "hm_hidden", ["m:head.heatmap.hidden"])
swap(graph, "hm_out", ["m:head.heatmap"])
swap_fn(tt_model, "sigmoid_heat", "m:head.sigmoid")
swap(graph, "local_max", ["m:head.local_max"])
qh = graph.query_head
swap_method(qh.select, "indices", ["m:query.topk"])
swap_method(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_method(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=str(SAMPLE))
ap.add_argument("--dispatch", default=None, choices=["eth", "worker"])
ap.add_argument("--cqs", type=int, default=None, choices=[1, 2])
ap.add_argument("--replays", type=int, default=1, help="traced frames inside the trace section")
ap.add_argument("--weights-dir", default=None)
ap.add_argument("--no-eager", dest="eager", action="store_false")
ap.add_argument("--no-layer-signposts", dest="layers", action="store_false")
ap.add_argument("--no-flush-modules", dest="flush_modules", action="store_false")
a = ap.parse_args()
import ttnn
from tt_bevfusion.tt.model import sparse_variant
t0 = time.perf_counter()
with BEVFusion.from_pretrained(dispatch=a.dispatch, num_command_queues=a.cqs,
weights_dir=a.weights_dir) as model:
tr, dev, g = model.runner, model.device, model.graph
print("loaded in %.1f s:" % (time.perf_counter() - t0), json.dumps(model.device_info), flush=True)
prepared = model._prepare(points=a.input)
out = model(a.input) # one served frame: upload + replays + readback
plan = prepared.plan
sv = sparse_variant(plan.bucket)
print(f"{out.to_dict()['num_detections']} detections, {prepared.num_voxels} voxels, bucket {plan.bucket}, "
f"counts {plan.counts()}", flush=True)
ttnn.synchronize_device(dev)
read_device_profiler(dev) # warm-up / capture / first frame out of the buffer
inputs = g.frame_inputs(prepared.vfe, plan, plan.bucket)
if a.eager:
def flush():
ttnn.synchronize_device(dev)
read_device_profiler(dev)
restore = module_signposts(g, layers=a.layers, flush=flush if a.flush_modules else None)
try:
with signposted("eager"):
tr.run_eager(sv, inputs=inputs)
ttnn.synchronize_device(dev)
flush()
tr.run_eager("dense")
ttnn.synchronize_device(dev)
finally:
restore()
read_device_profiler(dev)
tr.upload(inputs)
ttnn.synchronize_device(dev)
read_device_profiler(dev)
with signposted("trace"):
for _ in range(a.replays):
tr.replay(sv)
ttnn.synchronize_device(dev)
read_device_profiler(dev)
tr.replay("dense")
ttnn.synchronize_device(dev)
read_device_profiler(dev)
read_device_profiler(dev)
d = tr.describe()
print("trace:", json.dumps({"variant": sv, "program_cache_entries": d.get("program_cache_entries"),
"traces": d.get("traces"), "variants": d.get("variants")}), flush=True)
if __name__ == "__main__":
main()
|