# SPDX-License-Identifier: Apache-2.0 """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) # the bf16 input as float32 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()) # trace: 2 calls captured, replayed, equal to eager 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())