| |
| |
| """EPI-split bring-up probe (OPT_REPORT R2.3, ``tt/epi.py`` residual / row / split): the one-program forms vs the |
| stock op sequences they replace, bit for bit, per the hang protocol (smallest shape first; eager x2; 2-call trace). |
| |
| bin/devrun -t 600 -- env TT_METAL_WATCHER=2 python -u code/scripts/diag/epi_split_probe.py --smallest |
| bin/devrun -t 900 -- python -u code/scripts/diag/epi_split_probe.py |
| |
| Cases: |
| - ``proj``: the MBConv project conv epilogue: parts (1-3), fp32 bias, residual, split -> |
| stock ``finish(parts, bias, None, float32)`` then ``ttnn.add(., x)`` then ``split_activation(., 2)``; |
| - ``proj_nores``: the same without the residual (first block of a stage); |
| - ``gate``: the SE gate: part y, row s ``[1, 1, 1, C]`` -> stock ``split_activation(ttnn.multiply(y, s), 2)`` |
| (split only, no main output). |
| Prints one line per case and ``ALL OK`` / ``FAIL``. |
| """ |
| 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("--no-trace", action="store_true") |
| a = ap.parse_args() |
|
|
| import torch |
| import ttnn |
|
|
| from tt_sceneseg.device import close_device, open_device |
| from tt_sceneseg.tt import epi |
| from tt_sceneseg.tt.layers import _fp32_bias, finish, split_activation |
| from tt_sceneseg.ttaw.tensors import to_numpy |
|
|
| dev = open_device() |
| rng = np.random.default_rng(1) |
| shapes = [(32, 32), (32, 96)] if a.smallest else [(32, 32), (32, 96), (200, 1280 // 4), (12800, 24), (51200, 16), (3200, 40), |
| (800, 112), (12800, 144), (51200, 96), (200, 320), (224, 1152)] |
|
|
| def rand(r, c): |
| x = rng.standard_normal((1, 1, r, c)).astype(np.float32) * rng.choice([1e-3, 1.0, 30.0], size=(1, 1, r, 1)).astype(np.float32) |
| x[..., ::7, ::5] = 0.0 |
| return x |
|
|
| def up(x): |
| return ttnn.from_torch(torch.from_numpy(x), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev, |
| memory_config=ttnn.DRAM_MEMORY_CONFIG) |
|
|
| def bits(t, r, c): |
| v = to_numpy(t).reshape(-1, t.shape[-1])[:r, :c].astype(np.float32) |
| return v.view(np.uint32) |
|
|
| ok_all = True |
| t0 = time.time() |
| for (r, c) in shapes: |
| for case in ("proj", "proj_nores", "gate"): |
| nps = (1,) if case == "gate" else (1, 2, 3) |
| for np_ in nps: |
| xs = [rand(r, c) for _ in range(np_)] |
| res = rand(r, c) |
| b = (rng.standard_normal(c) * 0.5).astype(np.float32) |
| s = (1.0 / (1.0 + np.exp(-rng.standard_normal((1, 1, 1, c))))).astype(np.float32) |
| brow = _fp32_bias(b, dev) |
| btile = epi.bias_tile(b, dev) |
| |
| ps = [up(x) for x in xs] |
| if case == "gate": |
| y = ttnn.multiply(ps[0], up(s)) |
| ref_main = None |
| else: |
| y = finish(ps, brow, None, "float32") |
| if case == "proj": |
| y = ttnn.add(y, up(res)) |
| ref_main = bits(y, r, c) |
| ref_hi, ref_lo = (bits(t, r, c) for t in split_activation(y, 2)) |
| outs = [] |
| for _ in range(2): |
| pe = [up(x) for x in xs] |
| kw = dict(split=True) |
| if case == "gate": |
| kw.update(row=up(s), main=False) |
| bt = None |
| else: |
| bt = btile |
| if case == "proj": |
| kw.update(residual=up(res)) |
| if not epi.supported(pe, bt, None, residual=kw.get("residual"), row=kw.get("row")): |
| print(f"UNSUPPORTED {case} {(r, c)} np={np_}") |
| ok_all = False |
| break |
| m, hi, lo = epi.epilogue(pe, bt, None, "float32", **kw) |
| outs.append((None if m is None else bits(m, r, c), bits(hi, r, c), bits(lo, r, c))) |
| if len(outs) < 2: |
| continue |
| same = ((ref_main is None or np.array_equal(outs[0][0], ref_main)) and np.array_equal(outs[0][1], ref_hi) |
| and np.array_equal(outs[0][2], ref_lo)) |
| rep = all(np.array_equal(u, v) for u, v in zip(outs[0], outs[1]) if u is not None) |
| nd = [int((outs[0][i] != ref).sum()) for i, ref in ((0, ref_main), (1, ref_hi), (2, ref_lo)) |
| if ref is not None] |
| ok = same and rep |
| ok_all &= ok |
| print(f"{'OK ' if ok else 'FAIL'} {case:10s} {(r, c)} np={np_}: identical={same} (diff {nd}) " |
| f"eager_repeat={rep}", flush=True) |
|
|
| if not a.no_trace: |
| from tt_sceneseg.ttaw.trace import TraceRunner |
|
|
| r, c = (32, 32) if a.smallest else (12800, 144) |
| b = (rng.standard_normal(c) * 0.5).astype(np.float32) |
| btile = epi.bias_tile(b, dev) |
| brow = _fp32_bias(b, dev) |
| runner = TraceRunner(dev, name="epi_split_probe") |
| for nm in ("p0", "p1", "res", "y"): |
| runner.add_input(nm, rand(r, c), dtype="float32", layout=ttnn.TILE_LAYOUT) |
| runner.add_input("s", np.full((1, 1, 1, c), 0.5, np.float32), dtype="float32", layout=ttnn.TILE_LAYOUT) |
|
|
| def variant(ctx): |
| m, h1, l1 = epi.epilogue([ctx["p0"], ctx["p1"]], btile, None, "float32", residual=ctx["res"], split=True) |
| _, h2, l2 = epi.epilogue([ctx["y"]], None, None, "float32", row=ctx["s"], split=True, main=False) |
| return {"m": m, "h1": h1, "l1": l1, "h2": h2, "l2": l2} |
|
|
| runner.add_variant("v", variant) |
| runner.capture() |
| new = {nm: rand(r, c) for nm in ("p0", "p1", "res", "y")} |
| new["s"] = (1.0 / (1.0 + np.exp(-rng.standard_normal((1, 1, 1, c))))).astype(np.float32) |
| out = runner("v", inputs=new) |
| y1 = ttnn.add(finish([up(new["p0"]), up(new["p1"])], brow, None, "float32"), up(new["res"])) |
| h1, l1 = split_activation(y1, 2) |
| y2 = ttnn.multiply(up(new["y"]), up(new["s"])) |
| h2, l2 = split_activation(y2, 2) |
| refs = {"m": y1, "h1": h1, "l1": l1, "h2": h2, "l2": l2} |
| ok = True |
| for k, ref in refs.items(): |
| g = np.asarray(out[k]).reshape(-1, out[k].shape[-1])[:r, :c].astype(np.float32).view(np.uint32) |
| ok &= bool(np.array_equal(g, bits(ref, r, c))) |
| ok_all &= ok |
| print(f"{'OK ' if ok else 'FAIL'} trace 2-call {(r, c)}: all outputs identical={ok}", flush=True) |
| runner.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()) |
|
|
