| |
| """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 |
|
|
| |
| 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) |
| |
| 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) |
| |
| 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: |
| 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()) |
|
|
