File size: 9,067 Bytes
d160c10 | 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 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | # SPDX-License-Identifier: Apache-2.0
"""Probe (OPT round 4, item "L1 residency of layers 3-4"): one ResNet-50 identity bottleneck at a real encoder shape,
the round-3 path (DRAM between ops: ``Linear1x1`` conv1, ``L1Fallback`` conv2 + DB, stock 1x1 conv3, DRAM add) against
an L1-resident chain (residual stream block-sharded in L1; conv1 / conv3 as 2D-multicast ``ttnn.linear`` on the
sharded activations; conv2 reading conv1's shards; sharded add + ReLU). Random inputs / weights; reports bit identity
of every stage and the device time per block (trace of ``reps`` blocks, best of 3 replays).
bin/devrun -t 900 -- python -u code/scripts/probe_l1res.py --json logs/bevdet/opt_r4/probe_l1res.json
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
import numpy as np
# (name, batch, H, W, cin(=cout), mid)
SHAPES = {
"l3": (6, 16, 44, 1024, 256),
"l4": (6, 8, 22, 2048, 512),
}
def bits(t):
import torch
import ttnn
t = ttnn.to_memory_config(t, ttnn.DRAM_MEMORY_CONFIG) if t.memory_config().is_sharded() else t
return ttnn.to_torch(ttnn.from_device(t)).to(torch.bfloat16).view(torch.int16).numpy().reshape(-1, t.shape[-1])
def timed(device, fn, reps=10):
"""Device time per call (ms): ``fn`` returns tensors to free; one trace of ``reps`` calls, best of 3 replays."""
import ttnn
for t in fn():
ttnn.deallocate(t)
ttnn.synchronize_device(device)
tid = ttnn.begin_trace_capture(device, cq_id=0)
for _ in range(reps):
for t in fn():
ttnn.deallocate(t)
ttnn.end_trace_capture(device, tid, cq_id=0)
best = 1e9
for _ in range(3):
ttnn.synchronize_device(device)
t0 = time.perf_counter()
ttnn.execute_trace(device, tid, cq_id=0, blocking=True)
best = min(best, (time.perf_counter() - t0) * 1e3 / reps)
ttnn.release_trace(device, tid)
return best
def block_mc(rows, cols, gx, gy):
import ttnn
sh = -(-(-(-rows // 32)) // gy) * 32
grid = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))})
return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.BLOCK_SHARDED, ttnn.BufferType.L1,
ttnn.ShardSpec(grid, [sh, cols // gx], ttnn.ShardOrientation.ROW_MAJOR))
def pc2d(rows, k, n, kb, gx, gy, relu, sub=(2, 2)):
import ttnn
pm = -(-(-(-rows // 32)) // gy)
pn = n // 32 // gx
if pn <= 4:
sw = pn
sh = max(d for d in range(1, 5) if pm % d == 0 and d * sw <= 4)
else:
sh, sw = 1, max(d for d in range(1, 5) if pn % d == 0)
return ttnn.MatmulMultiCoreReuseMultiCastProgramConfig(
compute_with_storage_grid_size=(gx, gy), in0_block_w=kb, out_subblock_h=sh, out_subblock_w=sw,
out_block_h=pm, out_block_w=pn, per_core_M=pm, per_core_N=pn, transpose_mcast=False,
fused_activation=ttnn.UnaryWithParam(ttnn.UnaryOpType.RELU) if relu else None)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--json")
ap.add_argument("--only", default="l3")
ap.add_argument("--kb1", type=int, default=0, help="conv1 K block on the L1 chain (0: STOCK_KB)")
ap.add_argument("--kb3", type=int, default=0)
ap.add_argument("--gx", type=int, default=8)
a = ap.parse_args()
import torch
import ttnn
from tt_bevdet.device import close_device, open_device
from tt_bevdet.tt.backbone import STOCK_KB
from tt_bevdet.tt.conv_tune import DB
from tt_bevdet.tt.linear1x1 import wrap_1x1
from tt_bevdet.tt.params import DEFAULT_PRECISION
from tt_bevdet.tt.placement import L1Fallback
from tt_bevdet.ttaw.ops.conv import Conv2d, FeatureMap
from tt_bevdet.ttaw.precision import Precision
dev = open_device(trace_region_size=256 << 20)
g = dev.compute_with_storage_grid_size()
gy = int(g.y)
gx = a.gx
cfg = Precision.parse(DEFAULT_PRECISION).compute_kernel_config()
rng = np.random.default_rng(0)
res = []
relu = [ttnn.UnaryWithParam(ttnn.UnaryOpType.RELU)]
try:
for name in [s for s in a.only.split(",") if s]:
n, h, w, c, m = SHAPES[name]
rows = n * h * w
def mk(co, ci, k):
return ((rng.standard_normal((co, ci, k, k)) / np.sqrt(ci * k * k)).astype(np.float32),
(rng.standard_normal(co) * 0.1).astype(np.float32))
(w1, b1), (w2, b2), (w3, b3) = mk(m, c, 1), mk(m, m, 3), mk(c, m, 1)
x = np.maximum(rng.standard_normal((1, 1, rows, c)), 0).astype(np.float32)
xt = ttnn.from_torch(torch.from_numpy(x), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev,
memory_config=ttnn.DRAM_MEMORY_CONFIG)
# ---- reference: the round-3 path
c1 = wrap_1x1(Conv2d(w1, b1, padding=0, activation="relu", precision=DEFAULT_PRECISION, name="c1"))
c2 = L1Fallback(Conv2d(w2, b2, padding=1, activation="relu", precision=DEFAULT_PRECISION, name="c2"))
c2.conv.conv_config_overrides.update(DB)
c2.tune = dict(DB)
c3 = Conv2d(w3, b3, padding=0, precision=DEFAULT_PRECISION, name="c3")
def ref_block(keep=False):
fm = FeatureMap(xt, n, h, w, c)
y = c1(fm)
z = c2(y)
y3 = c3(z)
o = ttnn.add(y3.tensor, xt, activations=relu, memory_config=ttnn.DRAM_MEMORY_CONFIG)
if keep:
return [y.tensor, z.tensor, y3.tensor, o]
for t in (y, z, y3):
t.deallocate()
return [o]
ref = [bits(t) for t in ref_block(keep=True)]
t_ref = timed(dev, ref_block)
row = {"shape": name, "ref_ms": round(t_ref, 4), "c1": getattr(c1, "used", "conv2d"),
"c2": c2.placement}
print(f"[{name}] round-3 path: {t_ref:.4f} ms ({row['c1']}, conv2 {c2.placement})", flush=True)
# ---- L1 chain
kb1 = a.kb1 or STOCK_KB.get((rows, c, m), 1)
kb3 = a.kb3 or STOCK_KB.get((rows, m, c), 1)
xmc = block_mc(rows, c, gx, gy)
mmc = block_mc(rows, m, gx, gy)
xs = ttnn.to_memory_config(xt, xmc)
to = lambda a_: ttnn.from_torch(torch.from_numpy(np.ascontiguousarray(a_)), dtype=ttnn.float32,
layout=ttnn.TILE_LAYOUT, device=dev,
memory_config=ttnn.DRAM_MEMORY_CONFIG)
W1, B1 = to(w1[:, :, 0, 0].T), to(b1.reshape(1, -1))
W3, B3 = to(w3[:, :, 0, 0].T), to(b3.reshape(1, -1))
p1 = pc2d(rows, c, m, kb1, gx, gy, True)
p3 = pc2d(rows, m, c, kb3, gx, gy, False)
k2 = Conv2d(w2, b2, padding=1, activation="relu", precision=DEFAULT_PRECISION, name="k2",
output_memory_config=None, conv_config=dict(DB))
def l1_block(keep=False):
y = ttnn.linear(xs, W1, bias=B1, program_config=p1, compute_kernel_config=cfg, memory_config=mmc,
dtype=ttnn.bfloat16)
z = k2(FeatureMap(y, n, h, w, m)).tensor
if not keep:
ttnn.deallocate(y)
y3 = ttnn.linear(z, W3, bias=B3, program_config=p3, compute_kernel_config=cfg, memory_config=xmc,
dtype=ttnn.bfloat16)
if not keep:
ttnn.deallocate(z)
o = ttnn.add(y3, xs, activations=relu, memory_config=xmc)
if keep:
return [y, z, y3, o]
ttnn.deallocate(y3)
return [o]
try:
outs = l1_block(keep=True)
desc = [str(t.memory_config().shard_spec.shape) if t.memory_config().is_sharded() else "interleaved"
for t in outs]
got = [bits(t) for t in outs]
for t in outs:
ttnn.deallocate(t)
bad = [int(np.count_nonzero(gg[: rows] != rr[: rows])) for gg, rr in zip(got, ref)]
t_l1 = timed(dev, l1_block)
row.update(l1_ms=round(t_l1, 4), bad=bad, shards=desc, kb=[kb1, kb3])
print(f"[{name}] L1 chain kb {kb1}/{kb3}: {t_l1:.4f} ms; differing values y/z/y3/out {bad}; "
f"shards {desc}", flush=True)
except Exception as e: # noqa: BLE001
row["error"] = str(e)[:600]
print(f"[{name}] L1 chain error: {row['error']}", flush=True)
res.append(row)
for t in (xt, xs, W1, B1, W3, B3):
ttnn.deallocate(t)
for cc in (c1, c2, c3, k2):
cc.release()
finally:
close_device(dev)
if a.json:
Path(a.json).parent.mkdir(parents=True, exist_ok=True)
Path(a.json).write_text(json.dumps(res, indent=1))
return 0
if __name__ == "__main__":
sys.exit(main())
|