# SPDX-License-Identifier: Apache-2.0 """Probe (OPT round 2, item 7a): the image encoder's 1x1 / stride-1 convs as ``ttnn.linear`` with explicit program configs against ``ttnn.conv2d``'s own matmul (C17 ``Conv2d``), per shape: bit identity of the bf16 output and the device time of each candidate (trace replays of 20 calls, b2b). Candidates keep the stock K blocking (``in0_block_w``), so the fp32 K order is the same; ReLU is fused in the config (packer) where the conv has one. bin/devrun -t 1200 -- python -u code/scripts/probe_mm.py --json logs/bevdet/opt_r2/probe_mm.json """ from __future__ import annotations import argparse import itertools import json import sys import time from pathlib import Path import numpy as np # (rows, K, N, relu, stock kind, stock kb): the encoder's 1x1 stride-1 convs (OPT round 1 profile) SHAPES = [ (67584, 64, 64, True, "1d", 2), # layer1.0 conv1 (67584, 64, 256, False, "1d", 2), # layer1 conv3 / layer1.0 downsample (67584, 256, 64, True, "1d", 2), # layer1 conv1 (67584, 256, 128, True, "1d", 2), # layer2.0 conv1 (16896, 128, 512, False, "1d", 2), # layer2 conv3 (16896, 512, 128, True, "1d", 2), # layer2 conv1 (16896, 512, 256, True, "1d", 2), # layer3.0 conv1 (4224, 256, 1024, False, "2d", 1), # layer3 conv3 (4224, 1024, 256, True, "1d", 2), # layer3 conv1 (4224, 1024, 512, True, "2d", 1), # layer4.0 conv1 / lat0 (1056, 512, 2048, False, "2d", 1), # layer4 conv3 (1056, 2048, 512, True, "2d", 1), # layer4 conv1 / lat1 (16384, 256, 256, False, "1d", 2), # [12] decoder neck up2.4 (out_conv; OPT round 5) ] def divisors(n): return [d for d in range(1, n + 1) if n % d == 0] def subblocks(obh, obw): out = [] for sh, sw in [(2, 2), (1, 4), (4, 1), (1, 2), (2, 1), (1, 1)]: if obh % sh == 0 and obw % sw == 0: out.append((sh, sw)) return out[:1] def candidates(m, k, n, kb, gx, gy): import ttnn mt, nt, kt = m // 32, n // 32, k // 32 out = [] # 1D (in1 multicast, rows split over all cores) pcm = -(-mt // (gx * gy)) for obh in sorted({d for d in divisors(pcm) if d >= max(1, pcm // 6)} | {pcm}): for obw in sorted({d for d in divisors(nt)}): if obh * obw > 96 or obh * obw < 4 and not (obh == pcm and obw == nt): continue for sh, sw in subblocks(obh, obw): out.append(("1d", kb, obh, obw, sh, sw)) # 2D (block multicast) for cx in [d for d in divisors(nt) if d <= gx]: pm = -(-mt // gy) pn = nt // cx for obh in sorted({d for d in divisors(pm) if d >= max(1, pm // 8)}): if obh * pn > 96: continue for sh, sw in subblocks(obh, pn): out.append(("2d", kb, cx, obh, sh, sw)) return out def make_pc(c, m, n, gx, gy, relu): import ttnn mt, nt = m // 32, n // 32 act = ttnn.UnaryWithParam(ttnn.UnaryOpType.RELU) if relu else None if c[0] == "1d": _, kb, obh, obw, sh, sw = c return ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig( compute_with_storage_grid_size=(gx, gy), in0_block_w=kb, out_subblock_h=sh, out_subblock_w=sw, out_block_h=obh, out_block_w=obw, per_core_M=-(-mt // (gx * gy)), per_core_N=nt, fuse_batch=True, fused_activation=act, mcast_in0=False) _, kb, cx, obh, sh, sw = c pm = -(-mt // gy) return ttnn.MatmulMultiCoreReuseMultiCastProgramConfig( compute_with_storage_grid_size=(cx, gy), in0_block_w=kb, out_subblock_h=sh, out_subblock_w=sw, out_block_h=obh, out_block_w=nt // cx, per_core_M=pm, per_core_N=nt // cx, transpose_mcast=False, fused_activation=act) def bits(t): import torch import ttnn return ttnn.to_torch(ttnn.from_device(t)).to(torch.bfloat16).view(torch.int16).numpy() def timed(device, fn, reps=20): """Device time per call (ms): one trace of ``reps`` calls, replayed 3 times, best replay.""" import ttnn outs = [] fn() # compile ttnn.synchronize_device(device) tid = ttnn.begin_trace_capture(device, cq_id=0) for _ in range(reps): outs.append(fn()) 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) for o in outs: ttnn.deallocate(o) return best def main(): ap = argparse.ArgumentParser() ap.add_argument("--json") ap.add_argument("--only", default="", help="comma list of shape indices") a = ap.parse_args() import torch import ttnn from tt_bevdet.device import close_device, open_device from tt_bevdet.tt.params import DEFAULT_PRECISION from tt_bevdet.ttaw.ops.conv import Conv2d, FeatureMap from tt_bevdet.ttaw.precision import Precision dev = open_device(trace_region_size=512 << 20) g = dev.compute_with_storage_grid_size() gx, gy = int(g.x), int(g.y) prec = Precision.parse(DEFAULT_PRECISION) cfg = prec.compute_kernel_config() rng = np.random.default_rng(0) res = [] only = {int(v) for v in a.only.split(",") if v.strip()} try: for si, (m, k, n, relu, kind, kb) in enumerate(SHAPES): if only and si not in only: continue x = rng.standard_normal((1, 1, m, k)).astype(np.float32) w = (rng.standard_normal((n, k, 1, 1)) / np.sqrt(k)).astype(np.float32) b = rng.standard_normal(n).astype(np.float32) * 0.1 xt = ttnn.from_torch(torch.from_numpy(x), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) conv = Conv2d(w, b, padding=0, activation="relu" if relu else None, precision=DEFAULT_PRECISION, name=f"s{si}") fm = FeatureMap(xt, 1, m, 1, k) ref_t = conv(fm).tensor ref = bits(ref_t) ttnn.deallocate(ref_t) t_stock = timed(dev, lambda: conv(fm).tensor) wt = ttnn.from_torch(torch.from_numpy(np.ascontiguousarray(w[:, :, 0, 0].T)), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) bt = ttnn.from_torch(torch.from_numpy(b.reshape(1, -1)), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) row = {"shape": [m, k, n], "relu": relu, "stock_ms": round(t_stock, 4), "cands": []} print(f"[{si}] {m}x{k}x{n} relu={relu}: stock {t_stock:.4f} ms", flush=True) best = None for c in candidates(m, k, n, kb, gx, gy): try: pc = make_pc(c, m, n, gx, gy, relu) def run(): return ttnn.linear(xt, wt, bias=bt, program_config=pc, compute_kernel_config=cfg, memory_config=ttnn.DRAM_MEMORY_CONFIG, dtype=ttnn.bfloat16) o = run() got = bits(o) ttnn.deallocate(o) nbad = int(np.count_nonzero(got != ref)) t = timed(dev, run) except Exception as e: # noqa: BLE001 - config rejected / L1 clash row["cands"].append({"cfg": list(c), "error": str(e).splitlines()[0][:160]}) continue row["cands"].append({"cfg": list(c), "ms": round(t, 4), "bad": nbad}) print(f" {c}: {t:.4f} ms, {nbad} differ", flush=True) if nbad == 0 and (best is None or t < best[1]): best = (c, t) row["best_exact"] = None if best is None else {"cfg": list(best[0]), "ms": round(best[1], 4)} print(f" best exact: {row['best_exact']}", flush=True) res.append(row) for t in (xt, wt, bt): ttnn.deallocate(t) conv.release() finally: close_device(dev) if a.json: Path(a.json).write_text(json.dumps(res, indent=1)) return 0 if __name__ == "__main__": sys.exit(main())