| |
| """Probe (OPT round 6, R6.1): grid / shard orientation of the large block-sharded 3x3 convs at their model shapes. |
| |
| ttnn block-shards e.g. the eight 512 -> 512 3x3 convs of the CustomFPN / DepthNet ([6, 16, 44, 512]) ROW_MAJOR on |
| 8 x 10 cores (16 channel tiles do not split over 12 columns; 132 tile rows -> 14 per grid row). Candidates |
| ``T<x>x<y>``: ``transpose_shards`` (COL_MAJOR: rows over the x columns, channels over the y grid rows) with |
| ``override_sharding_config`` on an x-by-y grid; ``Tauto``: ttnn's own grid with ``transpose_shards``. Each candidate runs |
| (a) from a DRAM input with a DRAM output (``slicing="l1"``: I2S / S2I inside; the model's form for most convs) and |
| (b) from an input already in the candidate's own L1 shards with the output left there (the DN_L1 chain case). Per |
| candidate: output bits against the base (same input), max |diff|, device time per call (a trace of ``reps`` calls, |
| best of 3 blocking replays). |
| |
| bin/devrun -t 900 -- python -u code/scripts/probe_conv512.py --json logs/bevdet/opt_r6/probe_c512.json \ |
| --conv img_depth_net... --shape 6,16,44,512 --act relu --cands base,T12x8 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| def timed(device, fn, reps): |
| """Device time per call: a trace of ``reps`` calls, best of 3 blocking replays. The trace capture is always |
| ended (an error inside the capture would otherwise leave the device capturing).""" |
| import ttnn |
|
|
| fn().deallocate() |
| ttnn.synchronize_device(device) |
| outs = [] |
| tid = ttnn.begin_trace_capture(device, cq_id=0) |
| try: |
| for _ in range(reps): |
| outs.append(fn()) |
| finally: |
| ttnn.end_trace_capture(device, tid, cq_id=0) |
| best = 1e9 |
| try: |
| for _ in range(3): |
| ttnn.synchronize_device(device) |
| t0 = time.perf_counter() |
| ttnn.execute_trace(device, tid, cq_id=0, blocking=True) |
| best = min(best, (time.perf_counter() - t0) * 1e3 / reps) |
| finally: |
| ttnn.release_trace(device, tid) |
| for o in outs: |
| o.deallocate() |
| return best |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--json") |
| ap.add_argument("--reps", type=int, default=8) |
| ap.add_argument("--conv", default="img_bev_encoder_neck.conv.3", |
| help="module name for params.provider (DN-relative names start with 'dn:')") |
| ap.add_argument("--shape", default="6,16,44,512", help="n,h,w,cin") |
| ap.add_argument("--act", default="relu") |
| ap.add_argument("--cands", default="base,T12x8,T11x8,Tauto") |
| ap.add_argument("--abh", type=int, default=0, help="act_block_h_override for every candidate") |
| ap.add_argument("--nodb", action="store_true", help="no double buffering") |
| a = ap.parse_args() |
| import torch |
| import ttnn |
|
|
| from tt_bevdet.api import KNOBS |
| from tt_bevdet.device import close_device, open_device |
| from tt_bevdet.reference.config import default_config |
| from tt_bevdet.reference.weights import BEVDetWeights |
| from tt_bevdet.tests import paths |
| from tt_bevdet.tt.conv_tune import DB |
| from tt_bevdet.tt.params import DN, BEVDetParams |
| from tt_bevdet.ttaw.ops.conv import Conv2d, FeatureMap |
|
|
| cfg = default_config() |
| params = BEVDetParams(BEVDetWeights(paths.weights_onnx()), cfg) |
| prec = str(KNOBS.read().PRECISION) |
| name = DN + a.conv[3:] if a.conv.startswith("dn:") else a.conv |
| p = params.provider(name) |
| n, h, w, c = (int(v) for v in a.shape.split(",")) |
| BS = ttnn.TensorMemoryLayout.BLOCK_SHARDED |
| base = {} if a.nodb else dict(DB) |
| if a.abh: |
| base["act_block_h_override"] = a.abh |
| act = None if a.act in ("", "none") else a.act |
|
|
| def grid(x, y): |
| return ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(x - 1, y - 1))}) |
|
|
| def cand(label): |
| if label in ("base", "auto"): |
| return dict(base) |
| if label in ("autoT", "h4T", "w2T"): |
| return {**base, "shard_layout": BS, "transpose_shards": True} |
| if label == "autoBS": |
| return {**base, "shard_layout": BS} |
| if label == "h4": |
| return dict(base) |
| if label == "Tauto": |
| return {**base, "shard_layout": BS, "transpose_shards": True} |
| x, y = (int(v) for v in label[1:].split("x")) |
| return {**base, "shard_layout": BS, "override_sharding_config": True, "core_grid": grid(x, y), |
| "transpose_shards": label[0] == "T"} |
|
|
| SLICING = {"auto": "auto", "autoT": "auto", "autoBS": "auto", "h4": ("height", 4), "h4T": ("height", 4), |
| "w2T": ("width", 2)} |
| dev = open_device(trace_region_size=512 << 20) |
| print(f"grid {dev.compute_with_storage_grid_size()}, conv {name} {tuple(p.weight.shape)} in {a.shape} prec {prec}", |
| flush=True) |
| rng = np.random.default_rng(0) |
| rows = n * h * w |
| host = torch.from_numpy(rng.standard_normal((1, 1, rows, c)).astype(np.float32)).relu() |
| xin = ttnn.from_torch(host, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, |
| memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| fm = FeatureMap(xin, n, h, w, c) |
| ref = None |
| res = [] |
| try: |
| for label in [s for s in a.cands.split(",") if s]: |
| ov = cand(label) |
| cc = Conv2d(p.weight, p.bias, stride=p.stride, padding=p.padding, dilation=p.dilation, groups=p.groups, |
| activation=act, precision=prec, slicing=SLICING.get(label, "l1"), conv_config=ov, |
| name="probe." + label) |
| row = {"label": label} |
| try: |
| cc.output_memory_config = None if label not in SLICING else "dram" |
| o = cc(fm) |
| mc_out = o.tensor.memory_config() |
| got = ttnn.to_torch(ttnn.from_device(o.tensor)).float() |
| o.deallocate() |
| if ref is None: |
| ref = got |
| d = (got - ref).abs() |
| row.update(bad=int((d != 0).sum()), maxdiff=float(d.max()), |
| out_shard=(str(mc_out.shard_spec.shape) + " " + str(mc_out.shard_spec.grid.bounding_box())) |
| if mc_out.is_sharded() else "dram") |
| cc.output_memory_config = "dram" |
| row["ms_dram_io"] = round(timed(dev, lambda: cc(fm), a.reps), 4) |
| cc.output_memory_config = None |
| if mc_out.memory_layout == BS and label not in SLICING: |
| xs = ttnn.to_memory_config(xin, mc_out) |
| fs = FeatureMap(xs, n, h, w, c) |
| o2 = cc(fs) |
| row["chain_same_bits"] = bool((ttnn.to_torch(ttnn.from_device(o2.tensor)).float() == got).all()) |
| o2.deallocate() |
| row["ms_l1_chain"] = round(timed(dev, lambda: cc(fs), a.reps), 4) |
| ttnn.deallocate(xs) |
| except Exception as e: |
| row["error"] = str(e).splitlines()[0][:200] |
| finally: |
| cc.release() |
| dev.clear_program_cache() |
| print(json.dumps(row), flush=True) |
| res.append(row) |
| if a.json: |
| Path(a.json).write_text(json.dumps(res, indent=1)) |
| finally: |
| close_device(dev) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
