File size: 7,619 Bytes
e0adbab | 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 161 162 163 164 165 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Stem kernel probe (OPT_PLAN K-IN kernel form / MK0, ``tt/stem.py``), per the hang protocol.
``--smallest``: a 32 x 64 input (16 x 32 output, 16 tiles) with random uint8 values and the real stem weights:
eager x2 identical, a 2-call trace, and the error vs a float64 reference of the same math (conv of the uint8 values
with W_hi + W_lo, + bias map, SiLU). Default: the real 320 x 640 shape on the shipped sample, the kernel vs the stock
stem (``TtSceneSeg.stem`` with the knob off: bf16 cast, reshape, one conv per weight term, EPI), bit for bit, the
float64 error of both, and the time of each in a trace.
bin/devrun -t 600 -- env TT_METAL_WATCHER=2 python -u code/scripts/diag/stem_probe.py --smallest
bin/devrun -t 900 -- python -u code/scripts/diag/stem_probe.py
"""
from __future__ import annotations
import argparse
import sys
import time
import numpy as np
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--smallest", action="store_true")
ap.add_argument("--debug", action="store_true", help="the kernel's intermediate stages vs the host (smallest)")
a = ap.parse_args()
import torch
import torch.nn.functional as F
import ttnn
from tt_sceneseg.device import close_device, open_device
from tt_sceneseg.reference import config
from tt_sceneseg.reference.weights import load_weights
from tt_sceneseg.reference import assets
from tt_sceneseg.tt.layers import split_terms, stem_fold
from tt_sceneseg.tt.model import STEM_KEY
from tt_sceneseg.tt.stem import StemKernel
from tt_sceneseg.ttaw.tensors import to_numpy
weights = load_weights(assets.onnx_path())
lw = weights[STEM_KEY]
dev = open_device()
rc = 0
hw = (32, 64) if a.smallest else tuple(config.INPUT_HW)
w_u8, bias_rows = stem_fold(lw.weight, lw.bias, "bgr", hw)
rng = np.random.default_rng(1)
if a.smallest:
u8 = rng.integers(0, 256, size=(hw[0], hw[1], 3), dtype=np.uint8)
else:
from tt_sceneseg.host.preprocess import prepare_image
prep = prepare_image(str(config_sample()), channel_order="bgr")
u8 = np.asarray(prep.device_input(), np.uint8).reshape(hw[0], hw[1], 3)
img = ttnn.from_torch(torch.from_numpy(u8.reshape(1, 1, hw[0], hw[1] * 3).copy()), dtype=ttnn.uint8,
layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
bias = ttnn.from_torch(torch.from_numpy(bias_rows.reshape(1, 1, *bias_rows.shape).copy()), dtype=ttnn.float32,
layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
sk = StemKernel(w_u8, input_hw=hw)
# float64 reference of the same math
(wh, _), (wl, _) = split_terms(w_u8, np.zeros(32, np.float32), 2)
x = torch.from_numpy(u8.transpose(2, 0, 1)[None].astype(np.float64))
conv = F.conv2d(x, torch.from_numpy(wh.astype(np.float64) + wl.astype(np.float64)), stride=2, padding=1)
ref = F.silu(conv[0].permute(1, 2, 0).reshape(-1, 32) + torch.from_numpy(bias_rows.astype(np.float64))).numpy()
if a.debug:
A = np.zeros((hw[0] // 2 * hw[1] // 2, 32))
xp = np.pad(u8.astype(np.float64), ((1, 1), (1, 1), (0, 0)))
for oy in range(hw[0] // 2):
for ox in range(hw[1] // 2):
A[oy * (hw[1] // 2) + ox, :27] = xp[2 * oy:2 * oy + 3, 2 * ox:2 * ox + 3, :].reshape(27)
wt = sk.w_tiles.astype(np.float64)
exp = {1: A @ wt[:32], 2: bias_rows.astype(np.float64), 3: A @ wt[:32] + A @ wt[32:]}
for d in (1, 2, 3):
o = sk(img, bias, debug=d)
got = to_numpy(o).reshape(-1, 32).astype(np.float64)
ttnn.deallocate(o)
e = np.abs(got - exp[d]).max() / max(np.abs(exp[d]).max(), 1e-30)
bad = np.argwhere(np.abs(got - exp[d]) > 1e-3 * np.abs(exp[d]).max())
print(f"debug {d}: rel err {e:.2e}, {len(bad)} bad values, first {bad[:4].tolist()}", flush=True)
if len(bad):
r, c = bad[0]
print(" got", got[r, :8], "\n exp", exp[d][r, :8], flush=True)
outs = []
for _ in range(2):
o = sk(img, bias)
outs.append(to_numpy(o).reshape(-1, 32))
ttnn.deallocate(o)
same = np.array_equal(outs[0].view(np.uint32), outs[1].view(np.uint32))
err = float(np.abs(outs[0] - ref).max() / np.abs(ref).max())
tid = ttnn.begin_trace_capture(dev, cq_id=0)
o1 = sk(img, bias)
o2 = sk(img, bias)
ttnn.end_trace_capture(dev, tid, cq_id=0)
ttnn.execute_trace(dev, tid, cq_id=0, blocking=True)
tr = np.array_equal(to_numpy(o2).reshape(-1, 32).view(np.uint32), outs[0].view(np.uint32))
ttnn.release_trace(dev, tid)
for o in (o1, o2):
ttnn.deallocate(o)
# the FPU's in-tile accumulation of the 27 products is not exact fp32 (rel. err ~2e-4 of the output range on
# the smallest case, the A @ W_hi stage alone; the full run compares with the stock conv's own error)
ok = same and tr and err < 1e-3
line = f"stem kernel {hw}: eager x2 identical {same}, trace identical {tr}, rel err vs float64 {err:.2e}"
if not a.smallest:
from tt_sceneseg.tt.model import TtSceneSeg
# the served model's stem on the pixel-stick canvas (converted to row pages first) == the kernel on rows
mk = TtSceneSeg(weights, "bgr", stem_kernel=True)
pix = ttnn.from_torch(torch.from_numpy(u8.reshape(1, 1, hw[0] * hw[1], 3).copy()), dtype=ttnn.uint8,
layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
sp = mk.stem(pix)
pix_same = np.array_equal(to_numpy(sp.tensor).reshape(-1, 32).view(np.uint32), outs[0].view(np.uint32))
ttnn.deallocate(sp.tensor)
mk.release()
ok = ok and pix_same
line += f" | pixel-stick canvas identical {pix_same}"
model = TtSceneSeg(weights, "bgr", stem_kernel=False)
st = model.stem(img)
sn = to_numpy(st.tensor).reshape(-1, 32)
nd = int(np.sum(sn.view(np.uint32) != outs[0].view(np.uint32)))
serr = float(np.abs(sn - ref).max() / np.abs(ref).max())
ttnn.deallocate(st.tensor)
def tm(fn):
r = fn()
ttnn.deallocate(r)
tid = ttnn.begin_trace_capture(dev, cq_id=0)
rs = [fn() for _ in range(5)]
ttnn.end_trace_capture(dev, tid, cq_id=0)
ttnn.execute_trace(dev, tid, cq_id=0, blocking=True)
t1 = time.perf_counter()
for _ in range(5):
ttnn.execute_trace(dev, tid, cq_id=0, blocking=False)
ttnn.synchronize_device(dev)
us = (time.perf_counter() - t1) / 25 * 1e6
ttnn.release_trace(dev, tid)
for r in rs:
ttnn.deallocate(r)
return us
t_k = tm(lambda: sk(img, bias))
t_s = tm(lambda: model.stem(img).tensor)
ok = ok and err <= 1.5 * serr + 1e-7
line += (f" | vs stock stem: {nd} / {sn.size} values differ, stock rel err vs float64 {serr:.2e} | "
f"kernel {t_k:.1f} us, stock {t_s:.1f} us")
model.release()
print(line + (" OK" if ok else " FAIL"), flush=True)
rc |= 0 if ok else 1
sk.release()
close_device(dev)
print("ALL OK" if rc == 0 else "FAIL")
return rc
def config_sample():
from pathlib import Path
import tt_sceneseg
return Path(tt_sceneseg.__file__).parent / "samples" / "highway_normal_1.png"
if __name__ == "__main__":
sys.exit(main())
|