Download code/scripts/profile_ops.py from changh95/bevfusion-p150: direct link, hf CLI and curl.
- Browser
- Download file 9.06 kB
-
https://huggingface.co/changh95/bevfusion-p150/resolve/main/code/scripts/profile_ops.py
- Command line
-
hf download hf://changh95/bevfusion-p150/code/scripts/profile_ops.py
-
curl -L -o profile_ops.py https://huggingface.co/changh95/bevfusion-p150/resolve/main/code/scripts/profile_ops.py
9.06 kB
| #!/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() | |