#!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 """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) # stock 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())