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

"""Trace B alone -- the 4-step Stage-1 DiT denoise -- versus its DRAM byte floor (WP-06; IMPLEMENTATION_PLAN.md
§3.5, §5.3, §6; blackhole-perf-model-v2 §7.1; mb1 §0 / Table 1c).

What is reported

* ``denoise`` ms: median / p90 / min of ``--reps`` (50) replays of the denoise trace alone (``tracer.replay(denoise)``
  + ``synchronize``) after ``--warmup`` (10) untimed full calls (trace A must have run once so the hoisted K/V and
  ``state_features`` exist);
* the **bytes trace B streams** per inference under the dtype policy: every weight of the DiT blocks (self ``Wqkv``,
  cross ``Wq``, ``Wo``, ``ff1``, ``ff2`` + biases), the head (``proj_out_2``), the action encoder / decoder and the
  per-step AdaLN / ``b2_step{k}`` vectors, x 4 steps (the hoisted ``Wkv`` / ``bkv`` and the state encoder belong to
  trace A and are excluded).  Tile-padded device bytes come from ``tt.weights.upload_plan`` (WL §5 convention);
* ``bytes / 512 GB/s`` (GDDR6 spec, perf-model-v2 §2.1 / §7.1 floor) and ``bytes / 400 GB/s`` (the plan §5.3 arena
  budget = mb1 Table 1c bank-adjacent streaming rate for bfp8) in ms, the measured ratio to both and the effective
  GB/s, plus the plan §3.5 stock-op estimate from the ``MATMUL_CONFIGS`` table (measured matmul rows only).

``--dry-run`` (no device, no ttnn) does the byte accounting, resolves **every** trace-B matmul against
``tt.layers.MATMUL_CONFIGS`` at M = 64 (the DiT query rows) and every SDPA shape (self + the version's cross key
sets) against ``SDPA_CHUNKS`` -- any missing row raises -- and writes the JSON with ``dry_run: true``.  The device
path needs ``tt.model.Gr00tTT`` (WP-14; contract in ``bench_e2e``); without it, a non-dry run raises ``ImportError``.

Output ``benchmarks/results/bench_denoise_<version>_<ts>.json`` (schema ``gr00t-tt-perf/1``): ``stages.denoise``
(median ms) is what ``tests/tt/test_tt_perf.py`` compares; the byte model sits in ``meta.bytes``.

Run::

    DEVICE_LOCK_TIMEOUT=14400 /home/deepgadget/experiments/gr00t/bin/with-device.sh \\
        python -m models.experimental.gr00t.benchmarks.bench_denoise --version n16 --policy mixed_dit
    python -m models.experimental.gr00t.benchmarks.bench_denoise --version n15 --dry-run
"""

from __future__ import annotations

import argparse
import re
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple

from models.experimental.gr00t.benchmarks import bench_e2e as E2E
from models.experimental.gr00t.benchmarks import bench_stage_ops as SO
from models.experimental.gr00t.benchmarks import mb1_common as MB
from models.experimental.gr00t.tt import device as D
from models.experimental.gr00t.tt import layers as L

BENCH_NAME = "bench_denoise"
VERSIONS = SO.VERSIONS
DEFAULT_WARMUP = 10
DEFAULT_REPS = 50
GBS_SPEC = 512.0  # GDDR6 spec (perf-model-v2 §2.1)
GBS_ARENA = 400.0  # plan §5.3 arena budget / mb1 Table 1c bank-adjacent bfp8 rate
TRACE_B_CATEGORIES: Tuple[str, ...] = ("dit", "dit_precompute", "encoders", "encoders_slotted")
# trace-A tensors inside those categories (plan §3.3 hoist + state encoder; §4 trace table)
TRACE_A_NAME_RE = re.compile(r"^(dit\.block\.\d+\.(Wkv|bkv)|enc\.state\..*)$")
STEP_RE = re.compile(r"step(\d+)")
DIT_ROWS = 64  # plan D4: query rows padded to 64
# trace-B matmul weights whose (64, K, N, dtype) row must exist in MATMUL_CONFIGS (act = fused activation)
BLOCK_WEIGHTS: Tuple[Tuple[str, Optional[str]], ...] = (
    ("Wqkv", None),
    ("Wq", None),
    ("Wo", None),
    ("ff1.W", "gelu_tanh"),
    ("ff2.W", None),
)
OTHER_WEIGHTS: Tuple[Tuple[str, Optional[str]], ...] = (
    ("dit.proj_out_2.W", None),
    ("enc.action.W1", None),
    ("enc.action.W2a", None),
    ("enc.action.W3", None),
    ("enc.decoder.W1", "relu"),
    ("enc.decoder.W2", None),
)


# --------------------------------------------------------------------------------------------------------------------
# Byte model (CPU)
# --------------------------------------------------------------------------------------------------------------------
def trace_b_entries(version: str, embodiment_id: int, policy: Any) -> Dict[str, Any]:
    """``tt.weights.upload_plan`` entries that trace B streams (categories :data:`TRACE_B_CATEGORIES` minus the
    trace-A names :data:`TRACE_A_NAME_RE`)."""
    from models.experimental.gr00t.tt import weights as TW

    entries = TW.upload_plan(version, embodiment_id, policy, categories=TRACE_B_CATEGORIES)
    out = {n: e for n, e in entries.items() if not TRACE_A_NAME_RE.match(n)}
    if not out:
        raise KeyError(f"{version}: no trace-B tensors in categories {TRACE_B_CATEGORIES}")
    return out


def byte_model(entries: Mapping[str, Any], n_steps: int) -> Dict[str, Any]:
    """Bytes trace B reads per inference: ``n_steps x shared`` (every non-step tensor is streamed once per step) +
    the per-step tables (each read once).  Floors at :data:`GBS_SPEC` / :data:`GBS_ARENA`."""
    shared = 0.0
    per_step: Dict[int, float] = {k: 0.0 for k in range(n_steps)}
    by_dtype: Dict[str, float] = {}
    by_category: Dict[str, float] = {}
    for name, e in entries.items():
        b = float(e.nbytes)
        m = STEP_RE.search(name)
        if m:
            k = int(m.group(1))
            if k not in per_step:
                raise KeyError(f"{name}: step {k} outside the {n_steps} sampler steps")
            per_step[k] += b
        else:
            shared += b
        by_dtype[e.device_dtype] = by_dtype.get(e.device_dtype, 0.0) + b
        by_category[e.category] = by_category.get(e.category, 0.0) + b
    total = n_steps * shared + sum(per_step.values())
    return {
        "n_tensors": len(entries),
        "n_steps": n_steps,
        "shared_per_step_bytes": shared,
        "shared_per_step_mib": shared / (1 << 20),
        "step_tables_bytes": dict(sorted(per_step.items())),
        "total_bytes": total,
        "total_mib": total / (1 << 20),
        "by_dtype_mib": {k: v / (1 << 20) for k, v in sorted(by_dtype.items())},
        "by_category_mib": {k: v / (1 << 20) for k, v in sorted(by_category.items())},
        "floor_ms_512": total / (GBS_SPEC * 1e9) * 1e3,
        "floor_ms_400": total / (GBS_ARENA * 1e9) * 1e3,
    }


# --------------------------------------------------------------------------------------------------------------------
# Shape validation against the tables (CPU)
# --------------------------------------------------------------------------------------------------------------------
def _weight_shape(e: Any) -> Tuple[int, int]:
    s = tuple(int(x) for x in e.logical_shape)
    if len(s) != 2:
        raise ValueError(f"{e.name}: expected a 2-D weight, got {s}")
    return s


def validate_matmuls(cfg: Any, entries: Mapping[str, Any]) -> List[Dict[str, Any]]:
    """Every trace-B matmul at M = 64 must resolve in ``MATMUL_CONFIGS`` (plan §3.4 / §3.5); returns the plan rows
    with the table µs (``None`` for probes)."""
    rows: List[Dict[str, Any]] = []

    def add(name: str, act: Optional[str], where: str) -> None:
        e = entries[name]
        k, n = _weight_shape(e)
        dt = e.device_dtype
        if dt not in L.WEIGHT_DTYPES:
            raise ValueError(f"{name}: device dtype {dt!r} is not a matmul weight dtype {L.WEIGHT_DTYPES}")
        plan = L.matmul_plan(DIT_ROWS, k, n, dt, act)
        rows.append(
            {
                "where": where,
                "name": name,
                "K": k,
                "N": n,
                "dtype": dt,
                "act": act,
                "kind": plan.kind,
                "cores": plan.cores,
                "config": f"pcN={plan.per_core_N} ibw={plan.in0_block_w} sub={plan.out_subblock[0]}x{plan.out_subblock[1]}",
                "table_us": plan.spec.us,
                "probe": plan.spec.probe,
                "bytes": float(e.nbytes),
            }
        )

    for i in range(cfg.dit.n_blocks):
        pre = f"dit.block.{i}."
        is_cross = i % 2 == 0
        for stem, act in BLOCK_WEIGHTS:
            if stem == "Wqkv" and is_cross or stem == "Wq" and not is_cross:
                continue
            name = pre + stem
            if name not in entries:
                raise KeyError(f"plan tensor {name!r} missing from the trace-B entries")
            add(name, act, "cross" if is_cross else "self")
    for name, act in OTHER_WEIGHTS:
        if name not in entries:
            raise KeyError(f"plan tensor {name!r} missing from the trace-B entries")
        add(name, act, "head/encoders")
    return rows


def cross_sdpa_shapes(cfg: Any, layout: Any) -> List[Dict[str, Any]]:
    """``(class, Sq, Sk, d)`` of every DiT attention shape of the layout (plan §3.4 SDPA table) with its default
    chunks; ``KeyError`` when a shape has no allowlist row."""
    d = cfg.dit.head_dim_pad
    shapes: List[Dict[str, Any]] = [
        {"cls": "dit_self", "sq": DIT_ROWS, "sk": DIT_ROWS, "d": d, "mask": "const self mask"}
    ]
    key_sets = {cfg.dit.cross_key_subset(i) for i in cfg.dit.cross_block_indices}
    for ks in sorted(key_sets):
        if ks == "all":
            sk = layout.cross_all_keys[1]
            cls = "dit_cross_all_n15"
        elif ks == "text":
            sk = layout.cross_text_keys[1]
            cls = "dit_cross_text"
        elif ks == "image":
            sk = layout.cross_image_keys[1]
            cls = f"dit_cross_image_{cfg.version}"
        else:
            raise KeyError(f"unknown cross key set {ks!r}")
        shapes.append({"cls": cls, "sq": DIT_ROWS, "sk": int(sk), "d": d, "mask": f"{ks} keys"})
    for s in shapes:
        s["chunks"] = list(L.default_sdpa_chunks(s["cls"], s["sq"], s["sk"], s["d"]))
    return shapes


def table_estimate_us(cfg: Any, matmul_rows: Sequence[Mapping[str, Any]]) -> Dict[str, Any]:
    """Plan §3.5 stock-op estimate per step from the table: measured matmul µs summed over all blocks + head; probes
    counted, not estimated."""
    meas = sum(r["table_us"] for r in matmul_rows if r["table_us"] is not None)
    probes = [r["name"] for r in matmul_rows if r["table_us"] is None]
    return {
        "matmul_table_us_per_step": round(meas, 1),
        "n_matmuls": len(matmul_rows),
        "n_probes": len(probes),
        "probe_names_sample": probes[:6],
        "plan_3_5_step_ms": {"n15": 4.4, "n16": 8.6, "n17": 8.6}[cfg.version],
        "note": "plan §3.5 per-step estimate incl. head ops as measured by mb2; the table sum covers matmuls only",
    }


def analyse(version: str, policy: str, layout_name: Optional[str], embodiment: Optional[str]) -> Dict[str, Any]:
    """CPU analysis shared by the dry run and the device run."""
    # this benchmark measures the Stage-1 (ttnn) trace B -- its streamed bytes and its replay time; the megakernel
    # backend (the TTPolicy default since 2026-09-18) has no separate denoise op sequence to count
    setup = E2E.resolve_setup(
        version, policy=policy, layout_name=layout_name, embodiment=embodiment, dit_backend="ttnn"
    )
    cfg, layout = setup["cfg"], setup["layout"]
    n_steps = cfg.sampler.num_inference_timesteps
    entries = trace_b_entries(version, setup["embodiment_id"], setup["policy"])
    bytes_ = byte_model(entries, n_steps)
    mm_rows = validate_matmuls(cfg, entries)
    sdpa = cross_sdpa_shapes(cfg, layout)
    return {
        "setup": setup,
        "n_steps": n_steps,
        "bytes": bytes_,
        "matmuls": mm_rows,
        "sdpa": sdpa,
        "estimate": table_estimate_us(cfg, mm_rows),
    }


# --------------------------------------------------------------------------------------------------------------------
# Payload assembly
# --------------------------------------------------------------------------------------------------------------------
def _stage_from_ms(ms: Optional[Mapping[str, Any]], bytes_: Mapping[str, Any]) -> Dict[str, Any]:
    st: Dict[str, Any] = {
        "metric": "ms",
        "value": None,
        "floor_ms_512": round(bytes_["floor_ms_512"], 3),
        "floor_ms_400": round(bytes_["floor_ms_400"], 3),
        "total_mib": round(bytes_["total_mib"], 2),
    }
    if ms is not None:
        st.update(ms)
        v = float(ms["value"])
        st["x_floor_512"] = round(v / bytes_["floor_ms_512"], 3)
        st["x_floor_400"] = round(v / bytes_["floor_ms_400"], 3)
        st["effective_gbs"] = round(bytes_["total_bytes"] / (v * 1e-3) / 1e9, 1)
    return st


def build_payload(
    an: Mapping[str, Any],
    args: Mapping[str, Any],
    *,
    dry: bool,
    ms: Optional[Mapping[str, Any]] = None,
    device: Optional[Mapping[str, Any]] = None,
    extra_meta: Optional[Mapping[str, Any]] = None,
) -> Dict[str, Any]:
    setup = an["setup"]
    stages = {"denoise": _stage_from_ms(ms, an["bytes"])}
    meta = dict(setup["meta"])
    meta.update(
        bytes=an["bytes"],
        sdpa=an["sdpa"],
        estimate=an["estimate"],
        n_steps=an["n_steps"],
        trace_b_categories=list(TRACE_B_CATEGORIES),
    )
    if extra_meta:
        meta.update(extra_meta)
    return SO.perf_payload(
        BENCH_NAME,
        setup["cfg"].version,
        dry_run=dry,
        args=args,
        method={
            "denoise": "median / p90 / min ms of --reps replays of the denoise trace alone (replay + synchronize) after "
            "--warmup full calls (trace A run first); host wall clock",
            "bytes": "tt.weights.upload_plan tile-padded device bytes of the trace-B tensors x 4 steps (+ per-step tables once)",
            "floors": f"bytes / {GBS_SPEC:.0f} GB/s (GDDR6 spec) and / {GBS_ARENA:.0f} GB/s (plan §5.3 arena budget)",
        },
        stages=stages,
        rows=list(an["matmuls"]),
        device=device,
        meta=meta,
    )


def print_summary(payload: Mapping[str, Any]) -> None:
    m = payload["meta"]
    b = m["bytes"]
    st = payload["stages"]["denoise"]
    print(f"\n== {payload['bench']} {payload['version']} ({'DRY RUN' if payload['dry_run'] else 'measured'})")
    print(
        "   trace B streams {mib:.1f} MiB per inference ({n} tensors, {steps} steps; shared/step {sh:.1f} MiB; "
        "by dtype {dt}) -> floor {f512:.2f} ms @512 GB/s, {f400:.2f} ms @400 GB/s".format(
            mib=b["total_mib"],
            n=b["n_tensors"],
            steps=b["n_steps"],
            sh=b["shared_per_step_mib"],
            dt={k: round(v, 1) for k, v in b["by_dtype_mib"].items()},
            f512=b["floor_ms_512"],
            f400=b["floor_ms_400"],
        )
    )
    est = m["estimate"]
    print(
        f"   table estimate: {est['matmul_table_us_per_step']} µs/step of measured matmuls over {est['n_matmuls']} matmuls ({est['n_probes']} probes); plan §3.5 step ≈ {est['plan_3_5_step_ms']} ms"
    )
    print(
        "   sdpa shapes: "
        + "; ".join(f"{s['cls']} ({s['sq']},{s['sk']},{s['d']}) q{s['chunks'][0]}/k{s['chunks'][1]}" for s in m["sdpa"])
    )
    if st["value"] is not None:
        print(
            f"   denoise median {st['value']} ms (p90 {st['p90']}, min {st['min']}) = {st['x_floor_512']}x the 512 GB/s floor, {st['x_floor_400']}x the 400 GB/s floor; effective {st['effective_gbs']} GB/s"
        )
    else:
        print("   denoise ms: not measured (dry run)")


# --------------------------------------------------------------------------------------------------------------------
# Runs
# --------------------------------------------------------------------------------------------------------------------
def dry_run(
    *,
    version: str,
    policy: str = "mixed_dit",
    layout_name: Optional[str] = None,
    embodiment: Optional[str] = None,
    verbose: bool = True,
) -> Dict[str, Any]:
    an = analyse(version, policy, layout_name, embodiment)
    payload = build_payload(
        an,
        {
            "version": version,
            "policy": policy,
            "layout": an["setup"]["layout"].name,
            "embodiment": an["setup"]["embodiment_tag"],
        },
        dry=True,
        extra_meta={"gr00t_tt_importable": E2E.import_gr00t_tt() is not None},
    )
    if verbose:
        print_summary(payload)
    return payload


def run(
    device: Any,
    *,
    version: str,
    policy: str = "mixed_dit",
    trace_layout: str = "per_stage",
    sample: str = "canonical",
    layout_name: Optional[str] = None,
    embodiment: Optional[str] = None,
    warmup: int = DEFAULT_WARMUP,
    reps: int = DEFAULT_REPS,
    noise_seed: int = 0,
    cache_root: Optional[Path] = None,
    verbose: bool = True,
) -> Dict[str, Any]:
    """Time the denoise trace alone on an open device (model via ``bench_e2e.load_model``); returns the payload."""
    if warmup < 0 or reps <= 0:
        raise ValueError("warmup must be >= 0 and reps > 0")
    an = analyse(version, policy, layout_name, embodiment)
    setup = E2E.resolve_setup(
        version,
        policy=policy,
        trace_layout=trace_layout,
        layout_name=layout_name,
        embodiment=embodiment,
        dit_backend="ttnn",  # the Stage-1 denoise trace (see analyse); the device is the default-L1 one
    )
    ok, reason = E2E.golden_available(version, sample)
    if not ok:
        raise FileNotFoundError(f"{version}/{sample}: {reason}")
    inputs = E2E.load_golden_inputs(version, sample, setup["embodiment_tag"], setup["layout"].name)
    noise = E2E.make_noise(setup["layout"], noise_seed)
    model = E2E.load_model(setup, device, cache_root)
    tracer = E2E.resolve_stage_replayer(model, setup["stage_names"])
    stages = list(tracer.stage_names)
    den = E2E.denoise_stage_name(stages)
    tracer.upload(inputs["mi"], noise)
    tracer.synchronize()
    for _ in range(max(warmup, 1)):  # at least one full pass: trace A produces the hoisted K/V + state_features
        for st in stages:
            tracer.replay(st)
        tracer.synchronize()
    samples: List[float] = []
    for _ in range(reps):
        t0 = time.perf_counter()
        tracer.replay(den)
        tracer.synchronize()
        samples.append((time.perf_counter() - t0) * 1e3)
    ms = E2E._stats_ms(samples)
    facts = D.device_facts(device)
    facts.update(MB.device_report(__import__("ttnn"), device))
    payload = build_payload(
        an,
        {
            "version": version,
            "policy": policy,
            "trace_layout": trace_layout,
            "sample": sample,
            "layout": setup["layout"].name,
            "embodiment": setup["embodiment_tag"],
            "warmup": warmup,
            "reps": reps,
            "noise_seed": noise_seed,
        },
        dry=False,
        ms=ms,
        device=facts,
        extra_meta={
            "denoise_stage": den,
            "stage_names": stages,
            "shape_key": str(inputs["mi"].shape_key),
            "samples_ms": [round(x, 4) for x in samples],
        },
    )
    if verbose:
        print_summary(payload)
    return payload


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(
        description=__doc__.split("\n\n")[0], formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--version", required=True, choices=VERSIONS)
    ap.add_argument("--policy", default="mixed_dit", choices=E2E.POLICIES)
    ap.add_argument("--trace-layout", default="per_stage", choices=E2E.TRACE_LAYOUTS)
    ap.add_argument("--sample", default="canonical")
    ap.add_argument("--embodiment", default=None)
    ap.add_argument("--layout", default=None)
    ap.add_argument("--warmup", type=int, default=DEFAULT_WARMUP)
    ap.add_argument("--reps", type=int, default=DEFAULT_REPS)
    ap.add_argument("--noise-seed", type=int, default=0)
    ap.add_argument("--cache-root", default=None)
    ap.add_argument("--out", default=str(SO.RESULTS_DIR))
    ap.add_argument("--dry-run", action="store_true", help="byte model + table validation only; no device")
    ap.add_argument("--trace-region-mb", type=int, default=64)
    return ap


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = build_parser().parse_args(argv)
    if args.dry_run:
        payload = dry_run(version=args.version, policy=args.policy, layout_name=args.layout, embodiment=args.embodiment)
        SO.write_perf_json(payload, Path(args.out))
        return 0
    if E2E.import_gr00t_tt() is None:
        raise ImportError("tt.model.Gr00tTT is not importable (WP-14 pending) -- run with --dry-run")
    SO.ensure_tt_metal_cache()
    device = D.open_gr00t_device(trace_region_size=args.trace_region_mb << 20, l1_small_size=D.DEFAULT_L1_SMALL_SIZE)
    try:
        payload = run(
            device,
            version=args.version,
            policy=args.policy,
            trace_layout=args.trace_layout,
            sample=args.sample,
            layout_name=args.layout,
            embodiment=args.embodiment,
            warmup=args.warmup,
            reps=args.reps,
            noise_seed=args.noise_seed,
            cache_root=Path(args.cache_root) if args.cache_root else None,
        )
    finally:
        D.close_gr00t_device(device)
    SO.write_perf_json(payload, Path(args.out))
    return 0


__all__ = [
    "BENCH_NAME",
    "GBS_ARENA",
    "GBS_SPEC",
    "TRACE_B_CATEGORIES",
    "analyse",
    "byte_model",
    "cross_sdpa_shapes",
    "dry_run",
    "run",
    "trace_b_entries",
    "validate_matmuls",
]


if __name__ == "__main__":
    raise SystemExit(main())