File size: 6,517 Bytes
d160c10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | # 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())
|