#!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 """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 # (Hin, Win, C, k, s): the 16 SceneSeg depthwise layers (block_hw on the 320 x 640 input) 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) # compile + upload 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())