# 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())