File size: 6,015 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 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Round-6 probe (OPT_PLAN status after round 5, item 2): ttnn's DRAM-sliced decoder convs under explicit slice /
Conv2dConfig choices. Each decoder kxk conv that runs the DRAM-sliced path in the served graph (bf16 DRAM input, one
bf16 weight term, fp32 output, the decoder fidelity) is timed in a trace (10 calls, 5 replays) per configuration and
its output compared bit for bit with the automatic configuration.
bin/devrun -t 1500 -- python -u code/scripts/diag/slice_probe.py [--only dl5,z4] [--cfg auto,h2,...]
"""
from __future__ import annotations
import argparse
import sys
import time
import numpy as np
# name -> (H, W, Cin, Cout, k): the served graph's DRAM-sliced convs (input map size)
SHAPES = {"dl5": (80, 160, 512, 256, 3), "z4": (160, 320, 128, 512, 3), "dl7": (160, 320, 256, 128, 3),
"dl9": (320, 640, 128, 64, 3), "mrg": (320, 640, 64, 32, 3),
# L1-path convs (CONV_L1 = 3200 px): second 3x3s of blocks 0-1, S1 Z convs (k3 forms) of blocks 0-2,
# block 3's Z phases (80x160, DRAM-sliced)
"dl1": (20, 40, 768, 768, 3), "dl3": (40, 80, 512, 512, 3), "z0": (10, 20, 1280, 3072, 2),
"z1": (20, 40, 768, 2048, 2), "z2": (40, 80, 512, 2048, 2), "z3p": (80, 160, 256, 256, 2),
# backbone 1x1 convs (matmul path; stacked [W_hi ; W_lo] outputs): expand of stages 2 / 3 / 4 / 6,
# project of stages 2 / 3
"ex2": (160, 320, 16, 192, 1), "ex3": (80, 160, 24, 288, 1), "ex4": (40, 80, 40, 480, 1),
"ex6": (20, 40, 112, 1344, 1), "pr2": (80, 160, 96, 64, 1), "pr3": (40, 80, 144, 64, 1)}
def _configs():
import ttnn
hs, bs = ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.TensorMemoryLayout.BLOCK_SHARDED
db = {"enable_act_double_buffer": True, "enable_weights_double_buffer": True}
c = {"auto": ("auto", {}), "db": ("auto", db)}
for n in (2, 3, 4, 6, 8):
c[f"h{n}"] = (("height", n), {})
c[f"h{n}db"] = (("height", n), db)
for n in (2, 4, 8):
c[f"w{n}"] = (("width", n), {})
c["l1"] = ("l1", {})
c["l1db"] = ("l1", db)
c["l1adb"] = ("l1", {"enable_act_double_buffer": True})
c["l1wdb"] = ("l1", {"enable_weights_double_buffer": True})
c["adb"] = ("auto", {"enable_act_double_buffer": True})
c["wdb"] = ("auto", {"enable_weights_double_buffer": True})
for n in (2, 4):
c[f"h{n}hs"] = (("height", n), {"shard_layout": hs})
c[f"h{n}bs"] = (("height", n), {"shard_layout": bs})
for a in (64, 128, 256):
c[f"h{n}a{a}"] = (("height", n), {"act_block_h_override": a})
c["hs"] = ("auto", {"shard_layout": hs})
c["bs"] = ("auto", {"shard_layout": bs})
c["fid"] = ("auto", {"full_inner_dim": True})
return c
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--only", default="")
ap.add_argument("--cfg", default="")
ap.add_argument("--fid", default="HiFi3")
a = ap.parse_args()
import torch
import ttnn
from tt_sceneseg.device import close_device, open_device
from tt_sceneseg.tt.layers import _BiasFreeConv2d
from tt_sceneseg.ttaw.ops.conv import CONV_PRECISION, FeatureMap
from tt_sceneseg.ttaw.precision import Precision
from tt_sceneseg.ttaw.tensors import round_to_bf16, to_numpy
prec = Precision.parse(CONV_PRECISION).with_(fidelity=a.fid)
dev = open_device()
cfgs = _configs()
want = [c for c in a.cfg.split(",") if c] or list(cfgs)
base = "l1" if want[0].startswith("l1") else "auto"
if base not in want:
want = [base] + want
rng = np.random.default_rng(6)
for name, (H, W, ci, co, k) in SHAPES.items():
if a.only and name not in a.only.split(","):
continue
w = round_to_bf16((rng.standard_normal((co, ci, k, k)) / np.sqrt(ci * k * k)).astype(np.float32))
x = round_to_bf16(rng.standard_normal((H * W, ci)).astype(np.float32)).reshape(1, 1, H * W, ci)
xt = ttnn.from_torch(torch.from_numpy(x), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev,
memory_config=ttnn.DRAM_MEMORY_CONFIG)
fm = FeatureMap(xt, 1, H, W, ci)
ref = None
for cname in want:
slicing, cc = cfgs[cname]
conv = _BiasFreeConv2d(w, None, padding=k // 2, precision=prec, output_dtype="float32",
conv_config=cc, name=f"{name}.{cname}")
conv.slicing = slicing
try:
y = conv(fm).tensor
v = to_numpy(y).reshape(-1, y.shape[-1])[:H * W, :co].copy()
ttnn.deallocate(y)
if ref is None:
ref = v
same = "IDENTICAL" if np.array_equal(v.view(np.uint32), ref.view(np.uint32)) else \
f"DIFF max {float(np.abs(v - ref).max() / np.abs(ref).max()):.2e}"
tid = ttnn.begin_trace_capture(dev, cq_id=0)
for _ in range(10):
ttnn.deallocate(conv(fm).tensor)
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) / 50 * 1e6
ttnn.release_trace(dev, tid)
print(f"{name} {H}x{W} {ci}->{co} {cname:8s} {us:9.1f} us {same}", flush=True)
except Exception as e: # noqa: BLE001 (a rejected configuration is a result)
msg = str(e).strip().splitlines()
print(f"{name} {H}x{W} {ci}->{co} {cname:8s} FAILED {msg[0][:160] if msg else type(e).__name__}",
flush=True)
conv.release()
ttnn.deallocate(xt)
close_device(dev)
return 0
if __name__ == "__main__":
sys.exit(main())
|