# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""WP-K1 streaming gate G1: DRAM weight-stream rate of the DiT megakernel over the 32-block arena
(IMPLEMENTATION_PLAN.md §5.3, §5.7 K1 / G1: >= 384 GB/s (75 % of 512) in at least one reader mode for bf16, bfp8
recorded).

For every (stream dtype, reader mode) the script plans the arena of ``--version`` (all its DiT blocks, distinct random
weights), uploads it (WIDTH_SHARDED over the 8 banks), and times

* ``stream``: the four projections of every block with no gathers / mcasts (``DiTMegakernel.run_stream``) -- the pure
  weight-stream rate of the reader topology;
* ``chain``: the matmul-only block chain incl. the hub gathers / multicasts (``run_chain_timed``), with a PCC check of
  the last block's output against the torch chain.

Timing follows mb3 (device-stack-smoke.md §4.1): ``n`` in-kernel iterations, per-iteration time =
``(t(2n) - t(n)) / n`` (min of 3), GB/s = bytes streamed from DRAM per iteration / time.  With
``--profile`` (sets ``TT_METAL_DEVICE_PROFILER=1`` before ``ttnn`` is imported; one dtype x one mode per process, the
log is flushed at device close) a short run additionally collects the
``DeviceZoneScopedN`` zones (``mk_reader_op`` / ``mk_relay_op`` on the readers, ``mk_op`` / ``mk_P*`` on the compute
TRISCs, ``mk_hub_block``) from ``generated/profiler/.logs/profile_log_device.csv`` and reports, per launch, the
``dram_busy_fraction`` (relay: reader NCRISC ``mk_reader_op`` time / kernel span -- the readers are the only DRAM
masters) or the ``stream_active_fraction`` (direct: compute TRISC unpack-thread ``mk_op`` / ``mk_P*`` time / kernel
span -- a compute-active proxy, since the 96 NCRISC readers carry no zones).  Profile runs measure no rate, so their
JSON records ``g1.measured = false`` / ``g1.pass = null``.  ``--repeat N`` re-runs every (dtype, mode) measurement N
times in one process (a launch soak: ~17 launches per repeat; used by the WP-K1 fix round to bound the 22:18 hang).

Run through the device lock::

    DEVICE_LOCK_TIMEOUT=14400 bin/with-device.sh python models/experimental/gr00t/benchmarks/bench_mk_stream.py \\
        --version n16 --dtypes bf16,bfp8_b --modes relay,direct [--profile]

Results: ``benchmarks/results/bench_mk_stream_<version>_<timestamp>.json`` (schema ``gr00t-mk-k1/1``).
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import sys
import time
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence

HERE = Path(__file__).resolve().parent
RESULTS_DIR = HERE / "results"
TT_METAL_HOME = Path(os.environ.get("TT_METAL_HOME", str(HERE.parents[3])))
PROFILE_CSV = TT_METAL_HOME / "generated" / "profiler" / ".logs" / "profile_log_device.csv"
SCHEMA = "gr00t-mk-k1/1"
G1_GBPS = 384.0
PEAK_GBPS = 512.0
DEFAULT_ITERS = 4
PROFILE_BLOCKS = 4


def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    ap = argparse.ArgumentParser(
        description=__doc__.split("\n\n")[0], formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--version", default="n16", choices=("n15", "n16", "n17"))
    ap.add_argument("--dtypes", default="bf16,bfp8_b", help="comma-separated stream dtypes")
    ap.add_argument("--modes", default="relay,direct", help="comma-separated reader modes")
    ap.add_argument("--n-blocks", type=int, default=None, help="blocks of the arena (default: all of the version)")
    ap.add_argument("--iters", type=int, default=DEFAULT_ITERS, help="n of the (t(2n) - t(n)) / n timing")
    ap.add_argument("--ring-depth", type=int, default=None, help="DRAM pages in flight per PageRing")
    ap.add_argument("--relay-window", type=int, default=None, help="relay writes in flight per reader")
    ap.add_argument("--seed", type=int, default=11)
    ap.add_argument("--profile", action="store_true", help="also collect DeviceZoneScopedN zones (short run)")
    ap.add_argument("--skip-chain", action="store_true")
    ap.add_argument("--repeat", type=int, default=1, help="repeat every (dtype, mode) measurement N times (soak)")
    ap.add_argument("--out", default=str(RESULTS_DIR))
    return ap.parse_args(argv)


# --------------------------------------------------------------------------------------------------------------------
# profiler CSV
# --------------------------------------------------------------------------------------------------------------------
def read_zones(path: Path = PROFILE_CSV) -> Dict[str, Any]:
    """Aggregate the ``mk_*`` and ``*-KERNEL`` zones of the device profiler CSV **per launch** (``run host ID``
    column): per (RISC, zone) count / total / mean duration in microseconds and per-core totals, plus the kernel span
    (max over cores) of the launch."""
    if not path.is_file():
        raise FileNotFoundError(f"{path} not found; run with TT_METAL_DEVICE_PROFILER=1 (bench --profile)")
    with open(path) as fh:
        header = fh.readline()
        freq_mhz = float(header.split("CHIP_FREQ[MHz]:")[1].split(",")[0])
        reader = csv.reader(fh)
        cols = [c.strip() for c in next(reader)]
        idx = {c: i for i, c in enumerate(cols)}
        starts: Dict[tuple, int] = {}
        durations: Dict[tuple, List[int]] = defaultdict(list)
        for row in reader:
            if len(row) < len(cols) - 1:
                continue
            zone = row[idx["zone name"]].strip()
            if not (zone.startswith("mk_") or zone.endswith("-KERNEL")):
                continue
            run = row[idx["run host ID"]].strip()
            key = (run, int(row[idx["core_x"]]), int(row[idx["core_y"]]), row[idx["RISC processor type"]].strip(), zone)
            t = int(row[idx["time[cycles since reset]"]])
            typ = row[idx["type"]].strip()
            if typ == "ZONE_START":
                starts[key] = t
            elif typ == "ZONE_END" and key in starts:
                durations[key].append(t - starts.pop(key))
    runs: Dict[str, Dict[str, Any]] = {}
    for (run, x, y, risc, zone), ds in durations.items():
        r = runs.setdefault(run, {"zones": {}, "kernel_span_us": 0.0})
        name = f"{risc}:{zone}"
        agg = r["zones"].setdefault(name, {"count": 0, "total_us": 0.0, "cores": 0})
        agg["count"] += len(ds)
        agg["total_us"] += sum(ds) / freq_mhz
        agg["cores"] += 1
        if zone.endswith("-KERNEL"):
            r["kernel_span_us"] = max(r["kernel_span_us"], max(ds) / freq_mhz)
    for r in runs.values():
        for agg in r["zones"].values():
            agg["mean_us"] = agg["total_us"] / agg["count"] if agg["count"] else 0.0
            agg["per_core_total_us"] = agg["total_us"] / agg["cores"] if agg["cores"] else 0.0
    return {"freq_mhz": freq_mhz, "runs": runs}


def dram_busy_fraction(zones: Dict[str, Any], reader_mode: str) -> Dict[str, Optional[float]]:
    """Per launch: relay = reader NCRISC ``mk_reader_op`` time per core / kernel span; direct = compute TRISC_0
    (unpack thread) ``mk_op`` + ``mk_P*`` time per core / kernel span -- the fraction of the launch during which the
    weight stream is being consumed."""
    out: Dict[str, Optional[float]] = {}
    for run, r in zones["runs"].items():
        z = r["zones"]
        if reader_mode == "relay":
            busy = z.get("NCRISC:mk_reader_op", {}).get("per_core_total_us", 0.0)
        else:
            busy = sum(v["per_core_total_us"] for k, v in z.items() if k.startswith("TRISC_0:mk_"))
        span = r["kernel_span_us"]
        out[run] = busy / span if span > 0 and busy > 0 else None
    return out


# --------------------------------------------------------------------------------------------------------------------
# bench
# --------------------------------------------------------------------------------------------------------------------
def bench_one(ttnn: Any, device: Any, args: argparse.Namespace, dtype: str, reader_mode: str, profile: bool) -> Dict:
    import torch

    from models.experimental.gr00t.common import pcc as pcc_mod
    from models.experimental.gr00t.common.configs import get_config
    from models.experimental.gr00t.tt.megakernel import arena as A
    from models.experimental.gr00t.tt.megakernel import dit_program as P
    from models.experimental.gr00t.tt.megakernel.core_map import CoreMap

    n_blocks = args.n_blocks
    if profile and n_blocks is None:
        n_blocks = PROFILE_BLOCKS  # an explicit --n-blocks is honoured (zones beyond the L1 buffer are dropped)
    L = A.ArenaLayout.plan(
        args.version, n_blocks=n_blocks, dtype_stream=dtype, reader_mode=reader_mode, tile_order="row_major"
    )
    t0 = time.perf_counter()
    pack = L.pack(P.random_arena_plan(L, seed=args.seed))
    t_pack = time.perf_counter() - t0
    cm = CoreMap.from_device(device)
    kwargs: Dict[str, Any] = {}
    if args.ring_depth:
        kwargs["ring_depth"] = args.ring_depth
    if args.relay_window:
        kwargs["relay_window"] = args.relay_window
    mk = P.DiTMegakernel(get_config(args.version), None, None, L, cm, device, **kwargs)
    t0 = time.perf_counter()
    mk.upload(pack)
    ttnn.synchronize_device(device)
    t_upload = time.perf_counter() - t0
    n_blk = L.n_blocks
    bytes_iter = mk.stream_bytes(0, n_blk)
    x = torch.randn(P.M_ROWS, L.width, generator=torch.Generator().manual_seed(args.seed)).to(torch.bfloat16)
    x_rep = mk.make_x_rep(x)
    out: Dict[str, Any] = {
        "version": args.version,
        "dtype": dtype,
        "reader_mode": reader_mode,
        "n_blocks": n_blk,
        "bytes_per_iter": bytes_iter,
        "arena_bytes_per_bank": L.bytes_per_bank("stream"),
        "ring_depth": mk.ring_depth,
        "relay_window": mk.relay_window,
        "cb_w_slots": mk.ring_tiles // 16,
        "pack_s": t_pack,
        "upload_s": t_upload,
        "profile": profile,
    }
    try:
        n = 1 if profile else args.iters
        # ---- stream-only
        t0 = time.perf_counter()
        first = mk.run_stream(x_rep, 0, n_blk, 1)
        out["first_launch_stream_s"] = first
        out["compile_s"] = time.perf_counter() - t0
        if profile:
            mk.run_stream(x_rep, 0, n_blk, 1)
            out["stream"] = {"n_iters": 1}
        else:
            out["stream"] = P.measure_rate(lambda k: mk.run_stream(x_rep, 0, n_blk, k), n, 2 * n, bytes_iter)
            out["stream"]["gbps_frac_peak"] = out["stream"]["gbps"] / PEAK_GBPS
        # ---- chain
        if not args.skip_chain:
            got = mk.run_block(x_rep, 0, n_blk=n_blk, n_iters=1)
            ref = P.reference_block(L, pack, n_blk - 1, x.float())["g"]
            out["chain_pcc_last_block"] = pcc_mod.pcc(got.float(), ref)
            if profile:
                out["chain"] = {"n_iters": 1}
            else:
                out["chain"] = P.measure_rate(lambda k: mk.run_chain_timed(x_rep, 0, n_blk, k), n, 2 * n, bytes_iter)
                out["chain"]["gbps_frac_peak"] = out["chain"]["gbps"] / PEAK_GBPS
                out["chain"]["us_per_block"] = out["chain"]["per_iter_s"] / n_blk * 1e6
        if profile:
            ttnn.ReadDeviceProfiler(device)  # the CSV is flushed when the device closes: parsed in main()
        out["stats"] = mk.stats().__dict__
        out["kernel_groups"] = mk.kernel_groups()
        out["launches"] = mk.launches
    finally:
        ttnn.deallocate(x_rep)
        mk.deallocate()
    return out


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv)
    if args.profile:
        os.environ["TT_METAL_DEVICE_PROFILER"] = "1"
    os.environ.setdefault("TT_METAL_CACHE", str(Path.home() / ".cache" / "tt-metal-cache-gr00t"))
    import ttnn  # noqa: E402  (after the profiler env var)

    from models.experimental.gr00t.tt.device import close_gr00t_device, device_facts, open_gr00t_device

    dtypes = [d.strip() for d in args.dtypes.split(",") if d.strip()]
    modes = [m.strip() for m in args.modes.split(",") if m.strip()]
    if args.profile and (len(dtypes) != 1 or len(modes) != 1):
        raise SystemExit("--profile collects one device profiler log per process: pass exactly one dtype and one mode")
    if args.repeat < 1 or (args.profile and args.repeat != 1):
        raise SystemExit("--repeat must be >= 1 and is not combined with --profile (one profiler log per process)")
    if args.profile and PROFILE_CSV.exists():
        PROFILE_CSV.unlink()  # never attribute a previous run's zones to this one
    device = open_gr00t_device()
    runs: List[Dict[str, Any]] = []
    try:
        facts = device_facts(device)
        for dtype in dtypes:
            for mode in modes:
                for rep in range(args.repeat):
                    print(f"=== {args.version} {dtype} {mode} repeat {rep + 1}/{args.repeat} ===", flush=True)
                    r = bench_one(ttnn, device, args, dtype, mode, args.profile)
                    r["repeat"] = rep
                    runs.append(r)
                    s = r.get("stream", {})
                    c = r.get("chain", {})
                    head = f"RESULT {args.version} {dtype:7s} {mode:6s} blocks={r['n_blocks']} bytes/iter={r['bytes_per_iter']/1e6:.1f} MB "
                    if args.profile:
                        print(
                            head + f"profile run (no rate): {r['launches']} launches, "
                            f"chain PCC {r.get('chain_pcc_last_block', float('nan')):.5f}; zones parsed after close",
                            flush=True,
                        )
                    else:
                        print(
                            head
                            + f"stream={s.get('gbps', float('nan')):.0f} GB/s ({s.get('per_iter_s', float('nan'))*1e3:.2f} ms/iter) "
                            f"chain={c.get('gbps', float('nan')):.0f} GB/s ({c.get('us_per_block', float('nan')):.1f} us/block, "
                            f"PCC {r.get('chain_pcc_last_block', float('nan')):.5f}) launches={r['launches']}",
                            flush=True,
                        )
    finally:
        close_gr00t_device(device)
    if args.repeat > 1:
        pccs = [r["chain_pcc_last_block"] for r in runs if "chain_pcc_last_block" in r]
        print(
            f"SOAK {args.version}: {len(runs)} measurements, {sum(r['launches'] for r in runs)} launches completed, "
            f"chain PCC min {min(pccs) if pccs else float('nan'):.6f} max {max(pccs) if pccs else float('nan'):.6f}",
            flush=True,
        )
    busy_key = None
    if args.profile:
        zones = read_zones()
        busy = dram_busy_fraction(zones, modes[0])
        # relay: the readers are the DRAM masters -> DRAM-busy; direct: TRISC unpack-thread active time (proxy)
        busy_key = "dram_busy_fraction" if modes[0] == "relay" else "stream_active_fraction"
        runs[0]["zones"] = zones
        runs[0][busy_key] = busy
        # launches in order: stream (first, compiles), stream, chain
        labels = ["stream_cold", "stream", "chain"]
        for i, (run, r) in enumerate(sorted(zones["runs"].items(), key=lambda kv: int(kv[0]))):
            label = labels[i] if i < len(labels) else f"run{i}"
            print(
                f"PROFILE {args.version} {dtypes[0]} {modes[0]} {label}: kernel span {r['kernel_span_us']:.1f} us, "
                f"{busy_key} {busy.get(run)}"
            )
            for name, agg in sorted(r["zones"].items()):
                if name.endswith("-KERNEL"):
                    continue
                print(
                    f"ZONE {label:11s} {name:22s} count={agg['count']:5d} cores={agg['cores']:3d} "
                    f"mean_us={agg['mean_us']:8.2f} per_core_total_us={agg['per_core_total_us']:8.1f}"
                )
    measured = [r["stream"]["gbps"] for r in runs if "gbps" in r.get("stream", {}) and r["dtype"] == "bf16"]
    best = max(measured, default=0.0)
    doc = {
        "schema": SCHEMA,
        "bench": "bench_mk_stream",
        "version": args.version,
        "created": datetime.now().astimezone().isoformat(timespec="seconds"),
        "args": vars(args),
        "device": facts,
        "g1": {
            "gate_gbps": G1_GBPS,
            "measured": bool(measured),
            "best_bf16_stream_gbps": best if measured else None,
            "pass": (best >= G1_GBPS) if measured else None,
        },
        "busy_metric": busy_key,
        "runs": runs,
    }
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    tag = "profile_" if args.profile else ""
    path = out_dir / f"bench_mk_stream_{tag}{args.version}_{datetime.now():%Y%m%d-%H%M%S}.json"
    with open(path, "w") as fh:
        json.dump(doc, fh, indent=1, default=str)
        fh.write("\n")
    if measured:
        print(f"WROTE {path}  G1 bf16 best stream {best:.0f} GB/s -> {'PASS' if best >= G1_GBPS else 'FAIL'}")
    else:
        print(f"WROTE {path}  G1 not measured in this run (profile / non-bf16); see bench_mk_stream_n16_*.json")
    return 0


if __name__ == "__main__":
    sys.exit(main())