code/scripts/probe_mm.py

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