| |
| """Probe / bring-up (OPT round 5, R5.1 ``BEVDET_UPSAMPLE``): the TILE up-samplers of ``tt.upsample_k`` against their |
| host models, the float64 interpolation and the round-4 ``grid_sample`` path, on random bf16 maps; eager x2 equality, |
| a 2-call trace replay, and the device time per call (trace of ``reps`` calls, best of 3 replays). |
| |
| bin/devrun -t 600 -- env TT_METAL_WATCHER=2 python -u code/scripts/probe_upk.py --small # hang protocol step 1 |
| bin/devrun -t 900 -- python -u code/scripts/probe_upk.py --json logs/bevdet/opt_r5/probe_upk.json |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
| SMALL = [("upk", (8, 32), (16, 64), 32, 1), ("upk", (8, 32), (16, 64), 32, None), ("kron", (4, 4), (16, 16), 32, None), |
| ("upk", (4, 16), (16, 64), 32, 1), ("upk", (4, 16), (16, 64), 64, None)] |
| MODEL = [("upk", (64, 64), (128, 128), 512, None), ("upk", (16, 16), (64, 64), 640, None), |
| ("kron", (16, 16), (64, 64), 640, None)] |
|
|
|
|
| def to_host(t, h, w, c): |
| import torch |
| import ttnn |
|
|
| a = ttnn.to_torch(ttnn.from_device(t)).to(torch.float32).numpy().reshape(-1, t.shape[-1]) |
| return a[: h * w, :c].reshape(h, w, c) |
|
|
|
|
| def grid_path(dev, x_tile, h, w, c, H, W): |
| """The round-4 decoder path (tt.bev_decoder.TtBEVDecoder.upsample).""" |
| import torch |
| import ttnn |
|
|
| from tt_bevdet.ttaw.ops.deform import grid_sample, resize_grid |
|
|
| g = resize_grid((h, w), (H, W), align_corners=True) |
| grid = ttnn.from_torch(torch.from_numpy(g.copy()), dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, |
| memory_config=ttnn.DRAM_MEMORY_CONFIG) |
|
|
| def run(): |
| rm = ttnn.to_layout(x_tile, ttnn.ROW_MAJOR_LAYOUT, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| rm4 = ttnn.reshape(rm, (1, h, w, c)) |
| y = grid_sample(rm4, grid) |
| ttnn.deallocate(rm, False) |
| yf = ttnn.reshape(y, (1, 1, H * W, c)) |
| out = ttnn.to_layout(yf, ttnn.TILE_LAYOUT, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| ttnn.deallocate(y, False) |
| return out |
|
|
| return run |
|
|
|
|
| def timed(device, fn, reps=10): |
| import ttnn |
|
|
| ttnn.deallocate(fn()) |
| ttnn.synchronize_device(device) |
| tid = ttnn.begin_trace_capture(device, cq_id=0) |
| for _ in range(reps): |
| ttnn.deallocate(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) |
| return best |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--small", action="store_true") |
| ap.add_argument("--json", default="") |
| ap.add_argument("--terms", type=int, default=2) |
| ap.add_argument("--no-time", action="store_true") |
| a = ap.parse_args() |
| import torch |
| import ttnn |
|
|
| from tt_bevdet.device import close_device, open_device |
| from tt_bevdet.tt.upsample_k import KronUpsample, UpsampleK |
| from tt_bevdet.ttaw.ops.conv import FeatureMap |
| from tt_bevdet.ttaw.ops.upsample import interp_matrix |
|
|
| dev = open_device(trace_region_size=64 << 20) |
| res, ok = [], True |
| rng = np.random.default_rng(0) |
| try: |
| for kind, (h, w), (H, W), c, mc in (SMALL if a.small else MODEL): |
| x = rng.standard_normal((h, w, c)).astype(np.float32) * 2.0 |
| xt = ttnn.from_torch(torch.from_numpy(x.reshape(1, 1, h * w, c)), dtype=ttnn.bfloat16, |
| layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| xb = to_host(xt, h, w, c) |
| fm = FeatureMap(xt, 1, h, w, c) |
| if kind == "upk": |
| up = UpsampleK((h, w), (H, W), channels=c, terms=a.terms, max_cores=mc) |
| else: |
| up = KronUpsample((h, w), (H, W), channels=c, terms=a.terms) |
| o1 = up(fm) |
| r1 = to_host(o1.tensor, H, W, c) |
| o2 = up(fm) |
| r2 = to_host(o2.tensor, H, W, c) |
| o1.deallocate() |
| o2.deallocate() |
| eq = bool(np.array_equal(r1, r2)) |
| model = up.reference(xb) |
| exact = np.einsum("ph,hwc,qw->pqc", interp_matrix(h, H, mode="linear", align_corners=True), |
| xb.astype(np.float64), interp_matrix(w, W, mode="linear", align_corners=True)) |
| nd_model = int((r1 != model).sum()) |
| err_k = float(np.abs(r1 - exact).max()) |
| gp = grid_path(dev, xt, h, w, c, H, W) |
| g = gp() |
| rg = to_host(g, H, W, c) |
| ttnn.deallocate(g) |
| err_g = float(np.abs(rg - exact).max()) |
| bias_k = float((r1 - exact).mean()) |
| bias_g = float((rg - exact).mean()) |
| |
| ttnn.synchronize_device(dev) |
| tid = ttnn.begin_trace_capture(dev, cq_id=0) |
| t_a = up(fm) |
| t_b = up(fm) |
| ttnn.end_trace_capture(dev, tid, cq_id=0) |
| ttnn.execute_trace(dev, tid, cq_id=0, blocking=True) |
| tr_ok = bool(np.array_equal(to_host(t_a.tensor, H, W, c), r1) and |
| np.array_equal(to_host(t_b.tensor, H, W, c), r1)) |
| ttnn.release_trace(dev, tid) |
| t_a.deallocate() |
| t_b.deallocate() |
| row = {"kind": kind, "in": [h, w], "out": [H, W], "c": c, "max_cores": mc, "terms": a.terms, |
| "eager_equal": eq, "trace_equal": tr_ok, "neq_host_model": nd_model, "n": int(r1.size), |
| "max_err_kernel": err_k, "max_err_grid": err_g, "mean_err_kernel": bias_k, |
| "mean_err_grid": bias_g} |
| if not a.no_time and not a.small: |
| row["ms_kernel"] = timed(dev, lambda: up(fm).tensor) |
| row["ms_grid"] = timed(dev, gp) |
| ok &= eq and tr_ok |
| print(json.dumps(row), flush=True) |
| res.append(row) |
| up.release() |
| ttnn.deallocate(xt) |
| finally: |
| close_device(dev) |
| if a.json: |
| Path(a.json).write_text(json.dumps(res, indent=1)) |
| print("ALL OK" if ok else "FAILED") |
| return 0 if ok else 1 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
