| |
| """Probe (OPT round 3, R3.1): per-conv ``Conv2dConfig`` sweep of the model's kxk convs at their real shapes. |
| |
| The encoder and the decoder are built at the defaults and run once eagerly (random crop / gates / ``bev_in``) with |
| ``Conv2d.__call__`` hooked to record each kxk conv (its settled slicing and config overrides, its input spec). Every |
| distinct conv (by geometry + settled config) is then rebuilt with the same weights, bias, activation, precision and |
| output dtype and called on a random DRAM TILE input of that spec with candidate configs: activation / weight double |
| buffering, other shard layouts, activation block heights, ``full_inner_dim``, split reader / activation reuse. Per |
| candidate: the bf16 / fp32 output bit for bit against the settled config on the same input, and the device time per |
| call (one trace of ``reps`` calls, best of 3 blocking replays: the conv's whole call incl. its I2S / halo / S2I). |
| |
| bin/devrun -t 1500 -- python -u code/scripts/probe_conv.py --part enc --json logs/bevdet/opt_r3/probe_conv_enc.json |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| def bits(t): |
| import torch |
| import ttnn |
|
|
| x = ttnn.to_torch(ttnn.from_device(t)) |
| if x.dtype == torch.float32: |
| return x.view(torch.int32).numpy() |
| return x.to(torch.bfloat16).view(torch.int16).numpy() |
|
|
|
|
| def timed(device, fn, reps): |
| import ttnn |
|
|
| outs = [] |
| fn().deallocate() |
| ttnn.synchronize_device(device) |
| tid = ttnn.begin_trace_capture(device, cq_id=0) |
| for _ in range(reps): |
| outs.append(fn()) |
| ttnn.end_trace_capture(device, tid, cq_id=0) |
| best = 1e9 |
| 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) |
| ttnn.release_trace(device, tid) |
| for o in outs: |
| o.deallocate() |
| return best |
|
|
|
|
| def candidates(base_slicing, base_ov, cin, rows): |
| """(label, slicing, overrides) list; the first is the settled config.""" |
| import ttnn |
|
|
| HS, BS, WS = (ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.TensorMemoryLayout.BLOCK_SHARDED, |
| ttnn.TensorMemoryLayout.WIDTH_SHARDED) |
| b = dict(base_ov) |
| out = [("base", base_slicing, b)] |
| db = {"enable_act_double_buffer": True, "enable_weights_double_buffer": True} |
| out += [("adb", base_slicing, {**b, "enable_act_double_buffer": True}), |
| ("wdb", base_slicing, {**b, "enable_weights_double_buffer": True}), |
| ("db", base_slicing, {**b, **db})] |
| if base_slicing == "l1": |
| out += [("db+fid", "l1", {**b, **db, "full_inner_dim": True})] |
| else: |
| out += [(f"db+abh{h}", base_slicing, {**b, **db, "act_block_h_override": h}) for h in (256, 128, 64)] |
| out += [(f"hs+db{h or ''}", base_slicing, {**b, **db, "shard_layout": HS, **({"act_block_h_override": h} if h else {})}) |
| for h in (0, 128, 64)] |
| l1 = {k: v for k, v in b.items() if k != "act_block_h_override"} |
| for name, lay in (("hs", HS), ("bs", BS), ("ws", WS)): |
| for abh in (0, 128, 64): |
| if name == "ws" and abh: |
| continue |
| ov = {**l1, "shard_layout": lay} |
| if abh: |
| ov["act_block_h_override"] = abh |
| out.append((f"{name}{abh or ''}", "l1", ov)) |
| out.append((f"{name}{abh or ''}+db", "l1", {**ov, **db})) |
| if name == "hs": |
| out.append(("hs+db+reuse", "l1", {**l1, "shard_layout": HS, **db, "enable_activation_reuse": True})) |
| out.append(("hs+db+split", "l1", {**l1, "shard_layout": HS, **db, "force_split_reader": True})) |
| if name == "bs": |
| out.append(("bs+db+fid", "l1", {**l1, "shard_layout": BS, **db, "full_inner_dim": True})) |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--part", choices=("enc", "dec"), required=True) |
| ap.add_argument("--json") |
| ap.add_argument("--only", default="", help="comma list of conv-name substrings") |
| ap.add_argument("--reps", type=int, default=8) |
| ap.add_argument("--quick", action="store_true", help="only base / double-buffer candidates") |
| 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.params import BEVDetParams |
| from tt_bevdet.ttaw.ops import conv as convmod |
| from tt_bevdet.ttaw.ops.conv import Conv2d, FeatureMap |
|
|
| cfg = default_config() |
| params = BEVDetParams(BEVDetWeights(paths.weights_onnx()), cfg) |
| k = KNOBS.read() |
| prec = str(k.PRECISION) |
| rec = {} |
| orig = Conv2d.__call__ |
|
|
| def hook(self, x, **kw): |
| out = orig(self, x, **kw) |
| if not self.is_matmul and isinstance(x, FeatureMap): |
| t = x.tensor |
| rec[self.name] = {"conv": self, "spec": (x.batch, x.height, x.width, x.channels), |
| "slicing": self.slicing, "ov": dict(self.conv_config_overrides), |
| "dram_in": bool(convmod.is_dram(t)), "dtype": str(t.dtype)} |
| return out |
|
|
| dev = open_device(trace_region_size=1024 << 20) |
| rng = np.random.default_rng(0) |
| res = [] |
| try: |
| Conv2d.__call__ = hook |
| if a.part == "enc": |
| from tt_bevdet.tt.image_encoder import TtImageEncoder |
|
|
| mod = TtImageEncoder(params, precision=prec, stem_in="kin", conv_l1=str(k.CONV_L1), |
| backbone_res=str(k.BACKBONE_RES)) |
| n = cfg.n_cams |
| ih, iw = cfg.input_hw |
| crop = ttnn.from_torch(torch.from_numpy(rng.integers(0, 256, (1, 1, n * ih, iw * 3), dtype=np.uint8)), |
| dtype=ttnn.uint8, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, |
| memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| g = ttnn.from_torch(torch.full((n, 1, 1, 512), 0.5), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, |
| device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| d, f = mod(crop, g, g) |
| d.deallocate() |
| f.deallocate() |
| else: |
| from tt_bevdet.tt.bev_decoder import TtBEVDecoder |
| from tt_bevdet.tt.placement import parse_groups |
|
|
| dg = parse_groups(str(k.CONV_L1)) |
| mod = TtBEVDecoder(params, precision=prec, |
| conv_l1="all" if "decall" in dg else ("exact" if "dec" in dg else False), |
| merge=str(k.DEC_MERGE) not in ("0", "false", "False")) |
| ny, nx = cfg.bev_hw |
| x = ttnn.from_torch(torch.randn(1, 1, ny * nx, 864).relu(), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, |
| device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| mod(FeatureMap(x, 1, ny, nx, 864)).deallocate() |
| ttnn.deallocate(x) |
| Conv2d.__call__ = orig |
| mod.release() |
| seen = {} |
| for name, r in rec.items(): |
| if not r["dram_in"]: |
| continue |
| c = r["conv"] |
| key = (r["spec"], c.out_channels, c.kernel, c.stride, c.padding, c.dilation, c.activation, |
| c.output_dtype, r["slicing"], tuple(sorted(r["ov"].items()))) |
| seen.setdefault(key, []).append(name) |
| print(f"{len(rec)} kxk convs, {len(seen)} distinct", flush=True) |
| only = [s for s in a.only.split(",") if s] |
| for key, names in seen.items(): |
| if only and not any(o in nm for o in only for nm in names): |
| continue |
| r = rec[names[0]] |
| c = r["conv"] |
| bsz, h, w, cin = r["spec"] |
| rows = bsz * h * w |
| xin = ttnn.from_torch(torch.randn(1, 1, rows, cin).relu(), dtype=ttnn.bfloat16, |
| layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| fm = FeatureMap(xin, bsz, h, w, cin) |
| row = {"names": names, "spec": list(r["spec"]), "cout": c.out_channels, "kernel": list(c.kernel), |
| "stride": list(c.stride), "dilation": list(c.dilation), "slicing": r["slicing"], |
| "ov": {kk: str(v) for kk, v in r["ov"].items()}, "cands": []} |
| print(f"== {names} {r['spec']} -> {c.out_channels} k{c.kernel} s{c.stride} d{c.dilation} " |
| f"slicing={r['slicing']} ov={r['ov']}", flush=True) |
| ref = None |
| cands = candidates(r["slicing"], r["ov"], cin, rows) |
| if a.quick: |
| cands = cands[:4] |
| for label, slicing, ov in cands: |
| cc = Conv2d(c.weight, c.bias, stride=c.stride, padding=c.padding, dilation=c.dilation, |
| groups=c.groups, activation=c.activation, precision=c.precision, |
| output_dtype=c.output_dtype, slicing=slicing, conv_config=ov, name=c.name + "." + label) |
| try: |
| o = cc(fm) |
| got = bits(o.tensor) |
| o.deallocate() |
| if ref is None: |
| ref = got |
| nbad = int(np.count_nonzero(got != ref)) if got.shape == ref.shape else -1 |
| t = timed(dev, lambda: cc(fm), a.reps) |
| row["cands"].append({"label": label, "ms": round(t, 4), "bad": nbad}) |
| print(f" {label:14s} {t * 1e3:8.1f} us {nbad} differ", flush=True) |
| except Exception as e: |
| msg = str(e).splitlines()[0][:150] |
| row["cands"].append({"label": label, "error": msg}) |
| print(f" {label:14s} error: {msg}", flush=True) |
| if label == "base": |
| break |
| finally: |
| cc.release() |
| dev.clear_program_cache() |
| ok = [x for x in row["cands"] if "ms" in x and x["bad"] == 0] |
| if ok: |
| bst = min(ok, key=lambda x: x["ms"]) |
| row["best_exact"] = bst |
| print(f" best exact: {bst['label']} {bst['ms'] * 1e3:.1f} us (base " |
| f"{row['cands'][0].get('ms', 0) * 1e3:.1f}) x{len(names)}", flush=True) |
| res.append(row) |
| ttnn.deallocate(xin) |
| if a.json: |
| Path(a.json).write_text(json.dumps(res, indent=1)) |
| finally: |
| Conv2d.__call__ = orig |
| close_device(dev) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
