File size: 6,963 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 | #!/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())
|