| |
| |
| """K11 bring-up probe (OPT_PLAN K11, ``tt/dw.py``): the depthwise kernel vs a float64 CPU conv + bias + SiLU, per |
| the hang protocol (smallest shapes first; eager x2 bit-identical; a 2-call trace replayed on new inputs). |
| |
| bin/devrun -t 600 -- env TT_METAL_WATCHER=2 python -u code/scripts/diag/dw_probe.py --smallest |
| bin/devrun -t 900 -- python -u code/scripts/diag/dw_probe.py [--stock] |
| |
| Cases: random float32 maps (normal, with scaled rows and exact zeros) and EfficientNet-like weights. Per case: max |
| abs / rel error vs the fp64 reference (on the logical region) and whether two eager calls are bit-identical; |
| ``--stock`` also times nothing but reports the stock TermConv2d (two weight terms, two activation terms) error on the |
| same input for comparison. Prints one line per case and ``ALL OK`` / ``FAIL``. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| import time |
|
|
| import numpy as np |
|
|
| |
| GRAPH = [(160, 320, 32, 3, 1), (160, 320, 96, 3, 2), (80, 160, 144, 3, 1), (80, 160, 144, 5, 2), |
| (40, 80, 240, 5, 1), (40, 80, 240, 3, 2), (20, 40, 480, 3, 1), (20, 40, 480, 5, 1), |
| (20, 40, 672, 5, 1), (20, 40, 672, 5, 2), (10, 20, 1152, 5, 1), (10, 20, 1152, 3, 1)] |
| SMALL = [(8, 8, 32, 3, 1), (8, 8, 32, 3, 2), (8, 8, 32, 5, 1), (8, 8, 32, 5, 2), (5, 7, 40, 3, 1), (9, 6, 48, 5, 2)] |
|
|
|
|
| def ref(x, w, b, k, s, act): |
| import torch |
| import torch.nn.functional as F |
|
|
| C = w.shape[0] |
| xt = torch.from_numpy(x.astype(np.float64)) |
| y = F.conv2d(xt, torch.from_numpy(w.astype(np.float64)), torch.from_numpy(b.astype(np.float64)), stride=s, |
| padding=k // 2, groups=C) |
| if act == "silu": |
| y = F.silu(y) |
| return y.numpy() |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--smallest", action="store_true") |
| ap.add_argument("--no-trace", action="store_true") |
| ap.add_argument("--stock", action="store_true", help="also the stock TermConv2d error on the same inputs") |
| ap.add_argument("--bench", action="store_true", help="only time each graph layer (trace of 20 calls, 5 replays)") |
| ap.add_argument("--vs-base", action="store_true", |
| help="R7.1: also compare bit for bit with the DW_WPACK=0 form (one weight tile per tap)") |
| a = ap.parse_args() |
|
|
| import torch |
| import ttnn |
|
|
| from tt_sceneseg.device import close_device, open_device |
| from tt_sceneseg.tt.dw import DepthwiseKernel |
| from tt_sceneseg.ttaw.ops.conv import FeatureMap |
| from tt_sceneseg.ttaw.tensors import to_numpy |
|
|
| dev = open_device() |
| rng = np.random.default_rng(0) |
| cases = SMALL if a.smallest else SMALL + GRAPH |
|
|
| def mk(H, W, C): |
| x = rng.standard_normal((1, C, H, W)).astype(np.float32) |
| x *= rng.choice([1e-2, 1.0, 8.0], size=(1, 1, H, 1)).astype(np.float32) |
| x[..., ::7, ::5] = 0.0 |
| return x |
|
|
| def up(x): |
| C, H, W = x.shape[1:] |
| rows = np.ascontiguousarray(x[0].transpose(1, 2, 0).reshape(1, 1, H * W, C)) |
| return ttnn.from_torch(torch.from_numpy(rows), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev, |
| memory_config=ttnn.DRAM_MEMORY_CONFIG) |
|
|
| def down(t, Ho, Wo, C): |
| v = to_numpy(t).reshape(-1, t.shape[-1])[:Ho * Wo, :C] |
| return v.reshape(Ho, Wo, C).transpose(2, 0, 1)[None] |
|
|
| if a.bench: |
| tot = 0.0 |
| for (H, W, C, k, s) in GRAPH: |
| w = (rng.standard_normal((C, 1, k, k)) * 0.3).astype(np.float32) |
| b = (rng.standard_normal(C) * 0.5).astype(np.float32) |
| layer = DepthwiseKernel(w, b, stride=s, activation="silu", name="bench") |
| xt = up(mk(H, W, C)) |
| fm = FeatureMap(xt, 1, H, W, C) |
| ttnn.deallocate(layer(fm).tensor) |
| tid = ttnn.begin_trace_capture(dev, cq_id=0) |
| outs = [layer(fm).tensor for _ in range(20)] |
| 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) / 100 * 1e6 |
| tot += us |
| print(f"BENCH {(H, W, C, k, s)}: {us:8.1f} us per call", flush=True) |
| ttnn.release_trace(dev, tid) |
| for o in outs: |
| ttnn.deallocate(o) |
| ttnn.deallocate(xt) |
| layer.release() |
| print(f"BENCH total of the 12 shapes {tot:.1f} us", flush=True) |
| close_device(dev) |
| return 0 |
|
|
| ok_all = True |
| t0 = time.time() |
| for (H, W, C, k, s) in cases: |
| w = (rng.standard_normal((C, 1, k, k)) * 0.3).astype(np.float32) |
| b = (rng.standard_normal(C) * 0.5).astype(np.float32) |
| x = mk(H, W, C) |
| r = ref(x, w, b, k, s, "silu") |
| Ho, Wo = r.shape[2:] |
| layer = DepthwiseKernel(w, b, stride=s, activation="silu", name=f"dw{H}x{W}x{C}k{k}s{s}") |
| outs = [] |
| for _ in range(2): |
| xt = up(x) |
| y = layer(FeatureMap(xt, 1, H, W, C)) |
| outs.append(down(y.tensor, Ho, Wo, C)) |
| ttnn.deallocate(y.tensor) |
| ttnn.deallocate(xt) |
| rep = np.array_equal(outs[0].view(np.uint32), outs[1].view(np.uint32)) |
| if a.vs_base: |
| base = DepthwiseKernel(w, b, stride=s, activation="silu", name="base", wpack=False) |
| xt = up(x) |
| yb = base(FeatureMap(xt, 1, H, W, C)) |
| ob = down(yb.tensor, Ho, Wo, C) |
| ttnn.deallocate(yb.tensor) |
| ttnn.deallocate(xt) |
| base.release() |
| nb = int(np.sum(ob.view(np.uint32) != outs[0].view(np.uint32))) |
| print(f" vs base form: {nb} values differ (wpack {layer.wpack})", flush=True) |
| rep = rep and nb == 0 |
| err = np.abs(outs[0].astype(np.float64) - r) |
| rel = float(err.max() / max(np.abs(r).max(), 1e-30)) |
| ok = rep and np.all(np.isfinite(outs[0])) and rel < 1e-3 |
| extra = "" |
| if a.stock: |
| from tt_sceneseg.tt.layers import TermConv2d |
|
|
| st = TermConv2d(w, b, stride=s, padding=k // 2, groups=C, activation="silu", output_dtype="float32", |
| weight_terms=2, name="stock", epi=True) |
| xt = up(x) |
| ys = st(FeatureMap(xt, 1, H, W, C), act_terms=2) |
| es = np.abs(down(ys.tensor, Ho, Wo, C).astype(np.float64) - r) |
| extra = f" | stock max abs {es.max():.3e} mean {es.mean():.3e}" |
| ttnn.deallocate(ys.tensor) |
| st.release() |
| ok_all &= bool(ok) |
| print(f"{'OK ' if ok else 'FAIL'} {(H, W, C, k, s)} -> {Ho}x{Wo}: max abs {err.max():.3e} mean " |
| f"{err.mean():.3e} rel(max) {rel:.3e} eager_repeat={rep}{extra}", flush=True) |
| layer.release() |
|
|
| if not a.no_trace: |
| from tt_sceneseg.ttaw.trace import TraceRunner |
|
|
| tcases = [(8, 8, 32, 3, 1)] if a.smallest else [(8, 8, 32, 3, 2), (40, 80, 240, 5, 1)] |
| for (H, W, C, k, s) in tcases: |
| w = (rng.standard_normal((C, 1, k, k)) * 0.3).astype(np.float32) |
| b = (rng.standard_normal(C) * 0.5).astype(np.float32) |
| l1 = DepthwiseKernel(w, b, stride=s, activation="silu", name="t1") |
| l2 = DepthwiseKernel(w * 0.5, b, stride=1, activation=None, name="t2") |
| runner = TraceRunner(dev, name="dw_probe") |
| x0 = mk(H, W, C) |
| rows = np.ascontiguousarray(x0[0].transpose(1, 2, 0).reshape(1, 1, H * W, C)) |
| runner.add_input("x", rows, dtype="float32", layout=ttnn.TILE_LAYOUT) |
|
|
| def variant(ctx): |
| y1 = l1(FeatureMap(ctx["x"], 1, H, W, C)) |
| y2 = l2(y1) |
| return {"y1": y1.tensor, "y2": y2.tensor} |
|
|
| runner.add_variant("v", variant) |
| runner.capture() |
| x1 = mk(H, W, C) |
| rows1 = np.ascontiguousarray(x1[0].transpose(1, 2, 0).reshape(1, 1, H * W, C)) |
| out = runner("v", inputs={"x": rows1}) |
| r1 = ref(x1, w, b, k, s, "silu") |
| Ho, Wo = r1.shape[2:] |
| r2 = ref(r1.astype(np.float32), w * 0.5, b, k, 1, None) |
| g1 = np.asarray(out["y1"]).reshape(-1, out["y1"].shape[-1])[:Ho * Wo, :C].reshape(Ho, Wo, C).transpose(2, 0, 1)[None] |
| g2 = np.asarray(out["y2"]).reshape(-1, out["y2"].shape[-1])[:Ho * Wo, :C].reshape(Ho, Wo, C).transpose(2, 0, 1)[None] |
| e1 = float(np.abs(g1 - r1).max() / np.abs(r1).max()) |
| e2 = float(np.abs(g2 - r2).max() / np.abs(r2).max()) |
| ok = e1 < 1e-3 and e2 < 1e-3 |
| ok_all &= ok |
| print(f"{'OK ' if ok else 'FAIL'} trace 2-call {(H, W, C, k, s)}: rel {e1:.3e} / {e2:.3e}", flush=True) |
| runner.release() |
| l1.release() |
| l2.release() |
| print(f"{'ALL OK' if ok_all else 'FAIL'} in {time.time() - t0:.1f} s", flush=True) |
| close_device(dev) |
| return 0 if ok_all else 1 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
