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

"""End-to-end traced latency of ``Gr00tTT`` on one golden observation (WP-06; IMPLEMENTATION_PLAN.md §4, §6, §7.1).

Per call (plan §3.1): host ``encode`` -> upload of the per-call inputs -> trace replays (``per_stage``: vision, llm,
adapter, denoise; ``two``: A, B) -> D2H of ``action_pred`` -> host ``decode``.  After ``--warmup`` (10) untimed calls,
``--reps`` (50) timed calls give the **median / p90** of the end-to-end time and of every component, plus a second
pass that synchronises after every trace for the **per-trace split** (the split pass serialises the replays, so its
sum is >= the pipelined device time of the main pass; both are recorded).  Host encode / decode are timed on the
same observation.

Main pass = the production path ``tracer.call(mi, noise)`` (``TraceExecutor.replay``: static-shape asserts, host
inputs, writes -- on CQ 1 overlapping the traces when the device has two command queues -- and every trace, blocking
on the last one).  Its ``upload`` row is the executor's ``last_replay_upload_s`` (call start to the first
``execute_trace``: host preparation + the writes that gate the first trace), ``host_prepare`` the host-input part of
it, ``traces`` the rest of the call.  Split pass = ``tracer.upload`` (every input written on the trace CQ, serialised)
+ one synchronised replay per trace; its per-tensor accounting (``executor.upload_stats``: bytes, host tilize ms,
copy ms) is reported as ``upload/<input>`` rows and ``upload/total`` (docs/plan/n17-perf.md).

``--num-cqs`` (default 2) opens the device with that many command queues; ``--cq1-uploads auto|on|off`` and
``--embed-gather device|host`` are forwarded to ``Gr00tTT``; ``--profile-encode`` adds the per-step host-encode
profile (``meta.encode_profile``; also in ``--dry-run``).  ``--dit-backend ttnn|megakernel`` (WP-K5) selects the
denoise backend (``TTPolicy.dit_backend``; default = the policy default, the **megakernel** since 2026-09-18; the
device is opened through ``tt.model.open_model_device`` so the megakernel gets its ``worker_l1_size``) and
``--mk-arena-dtype auto|bf16|bfp8_b`` the megakernel's weight arena; ``meta.dit_backend`` / ``meta.mk_arena_dtype`` /
``meta.megakernel`` / ``meta.weights`` (tensors and bytes ``TTWeights`` uploaded, what the megakernel skipped) record
what ran.  ``--load-ttnn-dit-weights`` keeps the Stage-1 DiT upload under the megakernel (startup A/B).

The model comes from ``tt.model.Gr00tTT`` (WP-14).  Until that module exists the import is guarded: ``--dry-run``
validates ``--version`` / ``--policy`` / ``--trace-layout`` / ``--layout`` / ``--embodiment`` against ``common.configs``
and ``tt.policy.TTPolicy``, resolves the golden sample, times the **host** encode / decode when the golden observation
and the HF snapshot are present (CPU only, no ttnn), and writes the JSON with ``dry_run: true``; without ``--dry-run``
a missing ``tt.model`` raises ``ImportError`` (no silent fallback).

Contract expected from ``Gr00tTT`` (plan §2.2 rows ``tt/model.py`` / ``tt/tracer.py``; checked at run time by
:func:`resolve_stage_replayer`, which raises ``AttributeError`` naming every missing member):

* ``Gr00tTT.from_pretrained(version, embodiment=None, layout_name=None, policy=TTPolicy(), device=None,
  cache_root=None)``,
* ``model.tracer`` (a ``TraceExecutor``) with ``stage_names`` (execution order; ``("vision", "llm", "adapter",
  "denoise")`` or ``("A", "B")``), ``upload(mi, noise)`` (refresh the persistent per-call input buffers),
  ``replay(stage)`` (enqueue one trace, non-blocking), ``synchronize()`` and ``read_action_pred() -> torch.Tensor
  [1, 1, 64, A_pad]`` (D2H).

Output ``benchmarks/results/bench_e2e_<version>_<ts>.json`` (schema ``gr00t-tt-perf/1`` of ``bench_stage_ops``):
``stages`` = ``e2e`` (median ms, host included), ``device`` (upload -> D2H), ``trace/<stage>`` (split pass),
``upload``, ``d2h``, ``host_encode``, ``host_decode`` -- all ``metric: ms`` medians; p90 / min live next to them.

Run::

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

from __future__ import annotations

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

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.policy import DEFAULT_DIT_BACKEND, DIT_BACKENDS, MK_ARENA_DTYPES  # torch-free

BENCH_NAME = "bench_e2e"
VERSIONS = SO.VERSIONS
DEFAULT_WARMUP = 10
DEFAULT_REPS = 50
DEFAULT_HOST_REPS = 20
TRACE_LAYOUTS: Tuple[str, ...] = ("per_stage", "two")
POLICIES: Tuple[str, ...] = ("mixed_dit", "bf16", "mixed")
# plan §4: trace names per layout; the denoise trace is the last one of either layout
STAGE_NAMES: Dict[str, Tuple[str, ...]] = {"per_stage": ("vision", "llm", "adapter", "denoise"), "two": ("A", "B")}
REPLAYER_MEMBERS: Tuple[str, ...] = ("stage_names", "upload", "replay", "synchronize", "read_action_pred")
#: the production-path members the main pass needs on top of the split-pass contract (``tt.model.TraceReplayer``)
PRODUCTION_MEMBERS: Tuple[str, ...] = ("call", "upload_stats")
EMBED_GATHERS: Tuple[str, ...] = ("device", "host")
CQ1_MODES: Tuple[str, ...] = ("auto", "on", "off")


# --------------------------------------------------------------------------------------------------------------------
# Config / policy / golden resolution (CPU; shared with bench_denoise)
# --------------------------------------------------------------------------------------------------------------------
def resolve_setup(
    version: str,
    *,
    policy: str = "mixed_dit",
    trace_layout: str = "per_stage",
    layout_name: Optional[str] = None,
    embodiment: Optional[str] = None,
    dit_backend: str = DEFAULT_DIT_BACKEND,
    mk_arena_dtype: str = "auto",
) -> Dict[str, Any]:
    """Validate the CLI choices against ``common.configs`` / ``tt.policy`` and return the resolved facts
    (``cfg``, ``layout``, ``policy`` object, ``embodiment_tag`` / ``embodiment_id``, expected trace names)."""
    from models.experimental.gr00t.common.configs import get_config
    from models.experimental.gr00t.tt.policy import TTPolicy

    if version not in VERSIONS:
        raise ValueError(f"version must be one of {VERSIONS}, got {version!r}")
    if trace_layout not in TRACE_LAYOUTS:
        raise ValueError(f"trace_layout must be one of {TRACE_LAYOUTS}, got {trace_layout!r}")
    cfg = get_config(version)
    layout = cfg.canonical_layout if layout_name is None else cfg.layout(layout_name)
    tag = embodiment if embodiment is not None else layout.embodiment_tag
    emb_id = cfg.embodiment_id(tag)  # raises on an unknown tag
    pol = TTPolicy(
        dtype_policy=policy, trace_layout=trace_layout, dit_backend=dit_backend, mk_arena_dtype=mk_arena_dtype
    )  # raises on unknown values
    return {
        "cfg": cfg,
        "layout": layout,
        "policy": pol,
        "embodiment_tag": tag,
        "embodiment_id": emb_id,
        "stage_names": STAGE_NAMES[trace_layout],
        "meta": {
            "version": version,
            "layout_name": layout.name,
            "embodiment_tag": tag,
            "embodiment_id": emb_id,
            "s_pad_llm": layout.s_pad_llm,
            "s_pad_vl": layout.s_pad_vl,
            "n_images": layout.n_images,
            "dit_m_pad": layout.dit_m_pad,
            "noise_shape": list(layout.noise_shape),
            "policy": pol.to_dict(),
            "dit_backend": pol.dit_backend,
            "mk_arena_dtype": pol.mk_arena_dtype_for(version) if pol.dit_backend == "megakernel" else None,
            "expected_stage_names": list(STAGE_NAMES[trace_layout]),
        },
    }


def golden_available(version: str, sample: str = "canonical") -> Tuple[bool, str]:
    """``(ok, reason)`` -- the golden sample with its ``obs/`` and the pinned HF snapshot (tokenizer) are present."""
    from models.experimental.gr00t.common.configs import snapshot_dir
    from models.experimental.gr00t.common.golden import GoldenSet, golden_root

    if not golden_root().is_dir():
        return False, f"golden root {golden_root()} not present"
    try:
        gs = GoldenSet(version, sample)
    except (KeyError, FileNotFoundError) as e:
        return False, f"golden sample {sample!r} not present ({e})"
    from models.experimental.gr00t.reference.model import golden_obs_dir

    try:
        golden_obs_dir(gs)
    except FileNotFoundError as e:
        return False, str(e)
    try:
        snapshot_dir(version)
    except FileNotFoundError as e:
        return False, f"HF snapshot missing ({e})"
    return True, ""


def load_golden_inputs(
    version: str, sample: str, embodiment: Optional[str], layout_name: Optional[str]
) -> Dict[str, Any]:
    """``{gs, obs, mi}`` of a golden sample through ``reference.model.load_golden_sample`` + ``preprocessing.encode``
    (host only)."""
    from models.experimental.gr00t.common.preprocessing import encode
    from models.experimental.gr00t.reference.model import load_golden_sample

    gs, obs = load_golden_sample(version, sample)
    mi = encode(version, obs, embodiment=embodiment, layout_name=layout_name)
    return {"gs": gs, "obs": obs, "mi": mi}


def make_noise(layout: Any, seed: int) -> Any:
    """Deterministic ``initial_noise`` of the layout's ``noise_shape`` (``[1, H, A_max]``, fp32)."""
    import torch

    g = torch.Generator().manual_seed(int(seed))
    return torch.randn(tuple(layout.noise_shape), generator=g)


# --------------------------------------------------------------------------------------------------------------------
# Model loading and the stage-replayer contract (device)
# --------------------------------------------------------------------------------------------------------------------
def import_gr00t_tt() -> Any:
    """``tt.model.Gr00tTT`` or ``None`` when WP-14 has not landed (callers decide: dry-run vs ImportError)."""
    try:
        from models.experimental.gr00t.tt import model as M  # noqa: WPS433 - optional upstream module
    except ImportError:
        return None
    return getattr(M, "Gr00tTT", None)


def load_model(
    setup: Mapping[str, Any],
    device: Any,
    cache_root: Optional[Path] = None,
    *,
    embed_gather: str = "device",
    cq1_uploads: Optional[bool] = None,
    load_ttnn_dit_weights: Optional[bool] = None,
) -> Any:
    Gr00tTT = import_gr00t_tt()
    if Gr00tTT is None:
        raise ImportError(
            "tt.model.Gr00tTT is not importable (WP-14 pending); use --dry-run, or land tt/model.py with "
            "Gr00tTT.from_pretrained(version, embodiment, layout_name, policy, device, cache_root)"
        )
    if embed_gather not in EMBED_GATHERS:
        raise ValueError(f"embed_gather must be one of {EMBED_GATHERS}, got {embed_gather!r}")
    kw: Dict[str, Any] = dict(
        embodiment=setup["embodiment_tag"],
        layout_name=setup["layout"].name,
        policy=setup["policy"],
        device=device,
        embed_gather=embed_gather,
        cq1_uploads=cq1_uploads,
    )
    if load_ttnn_dit_weights is not None:
        kw["load_ttnn_dit_weights"] = bool(load_ttnn_dit_weights)
    if cache_root is not None:
        kw["cache_root"] = cache_root
    return Gr00tTT.from_pretrained(setup["cfg"].version, **kw)


def resolve_cq1(mode: str) -> Optional[bool]:
    """``"auto"`` -> ``None`` (``Gr00tTT`` enables the CQ-1 writes iff the device has two command queues),
    ``"on"`` / ``"off"`` -> ``True`` / ``False``."""
    if mode not in CQ1_MODES:
        raise ValueError(f"cq1_uploads must be one of {CQ1_MODES}, got {mode!r}")
    return None if mode == "auto" else mode == "on"


def resolve_stage_replayer(model: Any, expected_stages: Sequence[str]) -> Any:
    """``model.tracer`` satisfying the module-docstring contract; raises ``AttributeError`` listing missing members
    and ``ValueError`` when its ``stage_names`` are not the plan §4 names of the layout."""
    tracer = getattr(model, "tracer", None)
    if tracer is None:
        raise AttributeError("Gr00tTT has no `tracer` (plan §2.2 TraceExecutor); the benchmarks need it")
    missing = [m for m in REPLAYER_MEMBERS if not hasattr(tracer, m)]
    if missing:
        raise AttributeError(f"Gr00tTT.tracer lacks {missing}; required contract: {REPLAYER_MEMBERS}")
    names = tuple(str(s) for s in tracer.stage_names)
    if names != tuple(expected_stages):
        raise ValueError(f"tracer.stage_names {names} != expected {tuple(expected_stages)} for this trace layout")
    return tracer


def _worker_l1_size(device: Any) -> Optional[int]:
    """``tt.model.device_worker_l1_size`` (the allocatable L1 the device was opened with; recorded in ``meta``)."""
    try:
        from models.experimental.gr00t.tt.model import device_worker_l1_size
    except ImportError:
        return None
    return int(device_worker_l1_size(device))


def denoise_stage_name(stage_names: Sequence[str]) -> str:
    """The trace that holds the 4-step denoise: the last one of either layout (``denoise`` / ``B``)."""
    if not stage_names:
        raise ValueError("empty stage list")
    return str(stage_names[-1])


# --------------------------------------------------------------------------------------------------------------------
# Timing helpers
# --------------------------------------------------------------------------------------------------------------------
def _stats_ms(samples: Sequence[float]) -> Dict[str, Any]:
    if not samples:
        raise ValueError("no samples")
    xs = sorted(float(x) for x in samples)
    n = len(xs)
    p90 = xs[min(n - 1, int(round(0.9 * (n - 1))))]
    return {
        "metric": "ms",
        "value": round(statistics.median(xs), 4),
        "p90": round(p90, 4),
        "min": round(xs[0], 4),
        "max": round(xs[-1], 4),
        "mean": round(statistics.fmean(xs), 4),
        "n": n,
    }


def time_host(fn: Callable[[], Any], reps: int, warmup: int = 1) -> Dict[str, Any]:
    """Wall-clock ms of ``reps`` calls of ``fn`` after ``warmup`` untimed calls."""
    if reps <= 0:
        raise ValueError("reps must be positive")
    for _ in range(warmup):
        fn()
    out: List[float] = []
    for _ in range(reps):
        t0 = time.perf_counter()
        fn()
        out.append((time.perf_counter() - t0) * 1e3)
    return _stats_ms(out)


def host_timings(version: str, inputs: Mapping[str, Any], layout: Any, host_reps: int) -> Dict[str, Dict[str, Any]]:
    """Median / p90 ms of ``common.preprocessing.encode`` and ``common.normalization.decode`` (plan §3.1 host per
    call; f2 §1 measured 1.3 / 1.5 / 7.7 ms encode)."""
    import torch

    from models.experimental.gr00t.common.normalization import decode
    from models.experimental.gr00t.common.preprocessing import encode

    obs, mi = inputs["obs"], inputs["mi"]
    pred = torch.zeros(tuple(layout.noise_shape), dtype=torch.float32)
    return {
        "host_encode": time_host(
            lambda: encode(version, obs, embodiment=mi.embodiment_tag, layout_name=layout.name), host_reps
        ),
        "host_decode": time_host(lambda: decode(version, pred, obs, embodiment=mi.embodiment_tag), host_reps),
    }


def encode_profile(version: str, inputs: Mapping[str, Any], layout: Any, reps: int = 30) -> Dict[str, Dict[str, Any]]:
    """Per-step host-encode profile of one observation (median ms of ``reps`` calls after a warm-up): the version's
    image encoder and its parts, the prompt (cached and uncached tokenization), the state, the static layout, the
    RoPE / vision tables (cached and uncached) and the whole ``encode`` (docs/plan/n17-perf.md).  CPU only."""
    from models.experimental.gr00t.common import preprocessing as P
    from models.experimental.gr00t.common import prompts as PR
    from models.experimental.gr00t.common.configs import get_config
    from models.experimental.gr00t.common.normalization import encode_state, get_embodiment_spec

    obs, mi = inputs["obs"], inputs["mi"]
    cfg = get_config(version)
    spec = get_embodiment_spec(version, mi.embodiment_tag)
    frames = [obs.images[cam] for cam in spec.video_keys]
    encoders = {"n15": P.n15_encode_images, "n16": P.n16_encode_images, "n17": P.n17_encode_images}
    enc_images = encoders[version]
    vis = enc_images(frames, spec)
    tp = PR.encode_prompt(version, obs.instruction, vis.tokens_per_image)
    out: Dict[str, Dict[str, Any]] = {}
    out["images"] = time_host(lambda: enc_images(frames, spec), reps)
    if version == "n17":
        out["images/albumentations"] = time_host(
            lambda: [
                P.albumentations_eval_transform(
                    fr, spec.shortest_image_edge, spec.crop_fraction, spec.letter_box_transform
                )
                for f in frames
                for fr in f
            ],
            reps,
        )
        out["images/pixel_patches"] = time_host(lambda: P.qwen2vl_pixel_patches(vis.images_uint8), reps)
        out["images/transformers_processor"] = time_host(
            lambda: P.get_qwen2vl_image_processor()(images=[im for im in vis.images_uint8], return_tensors="pt"), reps
        )
    out["prompt/cached"] = time_host(lambda: PR.encode_prompt(version, obs.instruction, vis.tokens_per_image), reps)
    out["prompt/uncached"] = time_host(lambda: PR.tokenize_prompt(version, tp.text, vis.tokens_per_image), reps)
    out["state"] = time_host(lambda: encode_state(spec, obs.state), reps)
    lay = PR.select_layout(cfg, vis.n_images, vis.grid_thw, vis.tokens_per_image[0], tp.n_text, layout.name)
    out["slot_layout"] = time_host(lambda: PR.build_slot_layout(cfg, tp, lay), reps)
    slots = PR.build_slot_layout(cfg, tp, lay)
    inv_dtype = PR.LLM_ROPE_INV_FREQ_DTYPE[version]
    if version == "n17":
        grid = vis.image_grid_thw
        ids = slots.input_ids_padded[0]
        out["rope/cached"] = time_host(
            lambda: PR.mrope_tables_cached(
                tuple(ids.tolist()),
                cfg.prompt.image_token_id,
                PR.grid_key(grid),
                cfg.vision.merge_size,
                cfg.llm.head_dim,
                cfg.llm.rope_theta,
                tuple(cfg.llm.mrope_section),
                inv_dtype,
            ),
            reps,
        )

        def mrope_uncached() -> Any:
            pos3 = PR.mrope_position_ids(ids, cfg.prompt.image_token_id, grid, cfg.vision.merge_size)
            return PR.mrope_cos_sin(pos3, cfg.llm.head_dim, cfg.llm.rope_theta, cfg.llm.mrope_section, inv_dtype)

        out["rope/uncached"] = time_host(mrope_uncached, reps)
        out["vision_tables/cached"] = time_host(
            lambda: PR.vision_tables_cached(
                PR.grid_key(grid),
                cfg.vision.pos_grid,
                cfg.vision.merge_size,
                cfg.vision.head_dim,
                cfg.vision.rope_theta,
            ),
            reps,
        )
        out["vision_tables/uncached"] = time_host(
            lambda: (
                PR.vision_bilinear_indices_and_weights(grid, cfg.vision.pos_grid, cfg.vision.merge_size),
                PR.vision_rope_cos_sin(grid, cfg.vision.head_dim, cfg.vision.merge_size, cfg.vision.rope_theta),
                PR.vision_cu_seqlens(grid),
            ),
            reps,
        )
    else:
        out["rope"] = time_host(
            lambda: PR.rope_cos_sin_1d(slots.position_ids[0], cfg.llm.head_dim, cfg.llm.rope_theta, inv_dtype), reps
        )
    out["encode"] = time_host(
        lambda: P.encode(version, obs, embodiment=mi.embodiment_tag, layout_name=layout.name), reps
    )
    return out


# --------------------------------------------------------------------------------------------------------------------
# Device run
# --------------------------------------------------------------------------------------------------------------------
def run(
    device: Any,
    *,
    version: str,
    policy: str = "mixed_dit",
    trace_layout: str = "per_stage",
    sample: str = "canonical",
    embodiment: Optional[str] = None,
    layout_name: Optional[str] = None,
    warmup: int = DEFAULT_WARMUP,
    reps: int = DEFAULT_REPS,
    host_reps: int = DEFAULT_HOST_REPS,
    noise_seed: int = 0,
    cache_root: Optional[Path] = None,
    verbose: bool = True,
    embed_gather: str = "device",
    cq1_uploads: Optional[bool] = None,
    profile_encode: bool = False,
    dit_backend: str = DEFAULT_DIT_BACKEND,
    mk_arena_dtype: str = "auto",
    load_ttnn_dit_weights: Optional[bool] = None,
) -> Dict[str, Any]:
    """Measure the traced end-to-end latency and its split on an open device; returns the perf payload (not written).
    Main pass = production path (``tracer.call``), split pass = serialised writes + per-trace sync (module docstring).
    """
    import torch

    from models.experimental.gr00t.common.normalization import decode
    from models.experimental.gr00t.common.preprocessing import encode

    if warmup < 0 or reps <= 0:
        raise ValueError("warmup must be >= 0 and reps > 0")
    setup = resolve_setup(
        version,
        policy=policy,
        trace_layout=trace_layout,
        layout_name=layout_name,
        embodiment=embodiment,
        dit_backend=dit_backend,
        mk_arena_dtype=mk_arena_dtype,
    )
    ok, reason = golden_available(version, sample)
    if not ok:
        raise FileNotFoundError(f"{version}/{sample}: {reason}")
    inputs = load_golden_inputs(version, sample, setup["embodiment_tag"], setup["layout"].name)
    obs, mi = inputs["obs"], inputs["mi"]
    noise = make_noise(setup["layout"], noise_seed)
    t_load = time.perf_counter()
    model = load_model(
        setup,
        device,
        cache_root,
        embed_gather=embed_gather,
        cq1_uploads=cq1_uploads,
        load_ttnn_dit_weights=load_ttnn_dit_weights,
    )
    load_s = time.perf_counter() - t_load
    tracer = resolve_stage_replayer(model, setup["stage_names"])
    missing = [m for m in PRODUCTION_MEMBERS if not hasattr(tracer, m)]
    if missing:
        raise AttributeError(f"Gr00tTT.tracer lacks the production-path members {missing} (tt.model.TraceReplayer)")
    stages = list(tracer.stage_names)
    executor = model.executor
    per_tensor: List[Dict[str, Dict[str, float]]] = []

    def one_call(split: bool) -> Dict[str, float]:
        t: Dict[str, float] = {}
        t0 = time.perf_counter()
        mi_call = encode(version, obs, embodiment=setup["embodiment_tag"], layout_name=setup["layout"].name)
        t1 = time.perf_counter()
        if split:
            tracer.upload(mi_call, noise)
            tracer.synchronize()
            t2 = time.perf_counter()
            per_tensor.append(tracer.upload_stats)
            for st in stages:
                ts = time.perf_counter()
                tracer.replay(st)
                tracer.synchronize()
                t[f"trace/{st}"] = (time.perf_counter() - ts) * 1e3
            tracer.synchronize()
            t3 = time.perf_counter()
            t["host_prepare"] = model.timing["last_host_inputs_s"] * 1e3
        else:
            tracer.call(mi_call, noise)  # blocks on the last trace
            t3 = time.perf_counter()
            t2 = t1 + executor.timing["last_replay_upload_s"]
            t["host_prepare"] = executor.timing["last_replay_prepare_s"] * 1e3
        pred = tracer.read_action_pred()
        t4 = time.perf_counter()
        pred_norm = pred.float().reshape(pred.shape[-2], pred.shape[-1])[
            : setup["layout"].noise_shape[1], : setup["layout"].noise_shape[2]
        ]
        decode(version, pred_norm.unsqueeze(0), obs, embodiment=setup["embodiment_tag"])
        t5 = time.perf_counter()
        t.update(
            host_encode=(t1 - t0) * 1e3,
            upload=(t2 - t1) * 1e3,
            traces=(t3 - t2) * 1e3,
            d2h=(t4 - t3) * 1e3,
            host_decode=(t5 - t4) * 1e3,
            device=(t4 - t1) * 1e3,
            e2e=(t5 - t0) * 1e3,
        )
        return t

    for _ in range(warmup):
        one_call(split=False)
    main_pass = [one_call(split=False) for _ in range(reps)]
    split_pass = [one_call(split=True) for _ in range(reps)]
    keys_main = ("e2e", "device", "traces", "upload", "host_prepare", "d2h", "host_encode", "host_decode")
    summary: Dict[str, Dict[str, Any]] = {k: _stats_ms([c[k] for c in main_pass]) for k in keys_main}
    for st in stages:
        summary[f"trace/{st}"] = _stats_ms([c[f"trace/{st}"] for c in split_pass])
    summary["e2e_split_pass"] = _stats_ms([c["e2e"] for c in split_pass])
    summary["upload_split_pass"] = _stats_ms([c["upload"] for c in split_pass])
    summary["upload_split_pass"]["note"] = "tracer.upload (+sync): host inputs + every write serialised on the trace CQ"
    summary.update(upload_rows(per_tensor))
    summary["host_decode"]["note"] = "decode of the traced action_pred (valid rows / dims sliced on host)"
    summary["upload"]["note"] = "production path: call start -> first execute_trace (host inputs + gating writes)"
    facts = D.device_facts(device)
    facts.update(MB.device_report(__import__("ttnn"), device))
    ex_desc = executor.describe()
    meta = dict(
        setup["meta"],
        shape_key=str(mi.shape_key),
        n_text=mi.n_text,
        seq_len=mi.seq_len,
        model_load_s=round(load_s, 2),
        stage_names=stages,
        pred_shape=list(torch.as_tensor(tracer.read_action_pred()).shape),
        embed_gather=getattr(model, "embed_gather", None),
        cq1_uploads=ex_desc.get("cq1_uploads"),
        num_command_queues=int(executor.backend.num_command_queues(device)),
        prepare_asserts=ex_desc.get("prepare_asserts"),
        per_call_inputs=[s.describe() for s in executor.plan().inputs],
        per_call_bytes=ex_desc.get("plan", {}).get("per_call_bytes"),
        vit_fork=getattr(getattr(model.backbone, "tower", None), "fork", None),
        dit_backend=model.policy.dit_backend,
        mk_arena_dtype=model.recipe.mk_arena_dtype,
        megakernel=model.head.describe().get("megakernel"),
        worker_l1_size=_worker_l1_size(device),
        head_timing={k: v for k, v in model.head.timing.items()},
        model_timing={k: v for k, v in model.timing.items() if k.endswith("_s")},
        weights={
            "n_tensors": len(model.W),
            "device_bytes": float(model.W.device_bytes),
            "load_path": model.W.load_stats.get("path"),
            "load_s": model.timing.get("weights_load_s"),
            "skipped": dict(getattr(model, "weights_skipped", {})),
        },
    )
    if profile_encode:
        meta["encode_profile"] = encode_profile(version, inputs, setup["layout"])
    payload = SO.perf_payload(
        BENCH_NAME,
        version,
        dry_run=False,
        args={
            "version": version,
            "policy": policy,
            "trace_layout": trace_layout,
            "sample": sample,
            "embodiment": setup["embodiment_tag"],
            "layout": setup["layout"].name,
            "warmup": warmup,
            "reps": reps,
            "host_reps": host_reps,
            "noise_seed": noise_seed,
            "embed_gather": embed_gather,
            "cq1_uploads": cq1_uploads,
            "dit_backend": dit_backend,
            "mk_arena_dtype": mk_arena_dtype,
            "load_ttnn_dit_weights": load_ttnn_dit_weights,
        },
        method={
            "main_pass": f"{warmup} warm-up calls, then {reps} calls: encode -> tracer.call (TraceExecutor.replay: "
            "asserts, host inputs, writes, every trace; blocking on the last) -> D2H -> decode; medians / p90 of each "
            "component (device = call start..D2H; upload = call start..first execute_trace; traces = the rest)",
            "split_pass": f"{reps} more calls: tracer.upload (+sync) -> one synchronised replay per trace -> "
            "trace/<stage> medians and the per-tensor upload/<input> rows (bytes, host tilize ms, copy ms)",
            "clock": "host wall clock (time.perf_counter); trace replays enqueued non-blocking",
        },
        stages=summary,
        rows=[dict(c, pass_="main", i=i) for i, c in enumerate(main_pass)]
        + [dict(c, pass_="split", i=i) for i, c in enumerate(split_pass)],
        device=facts,
        meta=meta,
    )
    if verbose:
        print_summary(payload)
    return payload


def upload_rows(per_tensor: Sequence[Mapping[str, Mapping[str, float]]]) -> Dict[str, Dict[str, Any]]:
    """``upload/<input>`` stage rows (median ms of the serialised write; ``host_ms`` = from_torch / tilize,
    ``copy_ms`` = copy_host_to_device_tensor enqueue, ``bytes``) and ``upload/total`` from the split pass'
    ``executor.upload_stats`` snapshots."""
    if not per_tensor:
        return {}
    names = list(per_tensor[0])
    out: Dict[str, Dict[str, Any]] = {}
    for name in names:
        rows = [s[name] for s in per_tensor if name in s]
        st = _stats_ms([r["total_s"] * 1e3 for r in rows])
        st["host_ms"] = round(statistics.median([r["host_s"] * 1e3 for r in rows]), 4)
        st["copy_ms"] = round(statistics.median([r["copy_s"] * 1e3 for r in rows]), 4)
        st["bytes"] = int(rows[0]["bytes"])
        out[f"upload/{name}"] = st
    tot = _stats_ms([sum(r["total_s"] for r in s.values()) * 1e3 for s in per_tensor])
    tot["bytes"] = int(sum(int(s["bytes"]) for s in per_tensor[0].values()))
    tot["note"] = "sum of the serialised per-tensor writes (host tilize + copy enqueue) of one call"
    out["upload/total"] = tot
    return out


# --------------------------------------------------------------------------------------------------------------------
# Dry run
# --------------------------------------------------------------------------------------------------------------------
def dry_run(
    *,
    version: str,
    policy: str = "mixed_dit",
    trace_layout: str = "per_stage",
    sample: str = "canonical",
    embodiment: Optional[str] = None,
    layout_name: Optional[str] = None,
    host_reps: int = DEFAULT_HOST_REPS,
    verbose: bool = True,
    profile_encode: bool = False,
    dit_backend: str = DEFAULT_DIT_BACKEND,
    mk_arena_dtype: str = "auto",
) -> Dict[str, Any]:
    """Validate the setup and time the host side (encode / decode) when the golden data is present; no device."""
    setup = resolve_setup(
        version,
        policy=policy,
        trace_layout=trace_layout,
        layout_name=layout_name,
        embodiment=embodiment,
        dit_backend=dit_backend,
        mk_arena_dtype=mk_arena_dtype,
    )
    stages: Dict[str, Dict[str, Any]] = {}
    meta = dict(setup["meta"])
    meta["gr00t_tt_importable"] = import_gr00t_tt() is not None
    ok, reason = golden_available(version, sample)
    meta["golden_sample"] = sample
    if ok:
        inputs = load_golden_inputs(version, sample, setup["embodiment_tag"], setup["layout"].name)
        mi = inputs["mi"]
        meta.update(shape_key=str(mi.shape_key), n_text=mi.n_text, seq_len=mi.seq_len, n_images=mi.n_images)
        stages.update(host_timings(version, inputs, setup["layout"], host_reps))
        if profile_encode:
            prof = encode_profile(version, inputs, setup["layout"], reps=max(host_reps, 5))
            meta["encode_profile"] = prof
            stages.update({f"encode/{k}": v for k, v in prof.items()})
    else:
        meta["host_timing_skipped"] = reason
    for st in setup["stage_names"]:
        stages[f"trace/{st}"] = {"metric": "ms", "value": None, "note": "dry run: not measured"}
    payload = SO.perf_payload(
        BENCH_NAME,
        version,
        dry_run=True,
        args={
            "version": version,
            "policy": policy,
            "trace_layout": trace_layout,
            "sample": sample,
            "embodiment": setup["embodiment_tag"],
            "layout": setup["layout"].name,
            "host_reps": host_reps,
            "dit_backend": dit_backend,
            "mk_arena_dtype": mk_arena_dtype,
        },
        method={"host": f"{host_reps} host calls after 1 warm-up (encode / decode), median / p90 ms; no device"},
        stages=stages,
        rows=[],
        meta=meta,
    )
    if verbose:
        print_summary(payload)
    return payload


def print_summary(payload: Mapping[str, Any]) -> None:
    rows = []
    for k, s in payload["stages"].items():
        rows.append(
            {
                "stage": k,
                "median_ms": s.get("value"),
                "p90_ms": s.get("p90"),
                "min_ms": s.get("min"),
                "host_ms": s.get("host_ms"),
                "copy_ms": s.get("copy_ms"),
                "bytes": s.get("bytes"),
                "n": s.get("n"),
                "note": s.get("note"),
            }
        )
    MB.print_table(
        rows,
        ["stage", "median_ms", "p90_ms", "min_ms", "host_ms", "copy_ms", "bytes", "n", "note"],
        title=f"{payload['bench']} {payload['version']} ({'DRY RUN' if payload['dry_run'] else 'measured'})",
    )
    m = payload["meta"]
    print(
        "   layout {l} | embodiment {e} | S_pad {s} | shape_key {sk} | policy {p} | dit_backend {b} (arena {a}) | "
        "stages {st} | embed_gather {g} | cq1_uploads {c} ({n} CQ){skip}".format(
            l=m.get("layout_name"),
            e=m.get("embodiment_tag"),
            s=m.get("s_pad_llm"),
            sk=m.get("shape_key", "-"),
            p=(m.get("policy") or {}).get("dtype_policy"),
            b=m.get("dit_backend", (m.get("policy") or {}).get("dit_backend")),
            a=m.get("mk_arena_dtype", "-"),
            st=m.get("stage_names", m.get("expected_stage_names")),
            g=m.get("embed_gather", "-"),
            c=m.get("cq1_uploads", "-"),
            n=m.get("num_command_queues", "-"),
            skip=f" | host timing skipped: {m['host_timing_skipped']}" if m.get("host_timing_skipped") else "",
        )
    )


# --------------------------------------------------------------------------------------------------------------------
# CLI
# --------------------------------------------------------------------------------------------------------------------
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=POLICIES, help="TTPolicy.dtype_policy")
    ap.add_argument(
        "--dit-backend",
        default=DEFAULT_DIT_BACKEND,
        choices=DIT_BACKENDS,
        help=f"TTPolicy.dit_backend: Stage-1 ttnn ops | megakernel (default: the policy default, {DEFAULT_DIT_BACKEND})",
    )
    ap.add_argument(
        "--load-ttnn-dit-weights",
        action="store_true",
        help="keep the Stage-1 DiT / precompute / encoder upload under the megakernel backend (startup A/B; "
        "default: the megakernel skips them, tt.model.megakernel_weight_names)",
    )
    ap.add_argument(
        "--mk-arena-dtype",
        default="auto",
        choices=MK_ARENA_DTYPES,
        help="TTPolicy.mk_arena_dtype: the megakernel's weight arena (auto = the per-version default)",
    )
    ap.add_argument(
        "--trace-layout", default="per_stage", choices=TRACE_LAYOUTS, help="TTPolicy.trace_layout (plan §4)"
    )
    ap.add_argument("--sample", default="canonical", help="golden sample name")
    ap.add_argument("--embodiment", default=None, help="embodiment tag (default: the layout's)")
    ap.add_argument("--layout", default=None, help="static layout name (default: the version's canonical layout)")
    ap.add_argument("--warmup", type=int, default=DEFAULT_WARMUP)
    ap.add_argument("--reps", type=int, default=DEFAULT_REPS)
    ap.add_argument("--host-reps", type=int, default=DEFAULT_HOST_REPS)
    ap.add_argument("--noise-seed", type=int, default=0)
    ap.add_argument("--cache-root", default=None, help="TTWeights cache root (default: the plan cache)")
    ap.add_argument("--out", default=str(SO.RESULTS_DIR))
    ap.add_argument("--dry-run", action="store_true", help="validate the setup, time the host side only; no device")
    ap.add_argument("--trace-region-mb", type=int, default=64)
    ap.add_argument("--num-cqs", type=int, default=2, choices=(1, 2), help="command queues the device is opened with")
    ap.add_argument("--cq1-uploads", default="auto", choices=CQ1_MODES, help="per-call input writes on CQ 1")
    ap.add_argument("--embed-gather", default="device", choices=EMBED_GATHERS, help="Gr00tTT text-row gather")
    ap.add_argument("--profile-encode", action="store_true", help="add the per-step host-encode profile to meta")
    ap.add_argument(
        "--torch-threads", type=int, default=None, help="torch.set_num_threads for the host side (experiment knob)"
    )
    ap.add_argument(
        "--cv2-threads", type=int, default=None, help="cv2.setNumThreads for the host side (experiment knob)"
    )
    return ap


def set_host_threads(torch_threads: Optional[int], cv2_threads: Optional[int]) -> Dict[str, int]:
    """Process-level thread counts of the host libraries (n17-perf.md open issue: 16 torch + 32 cv2 threads make
    the small host ops 1.5-2x slower next to the ttnn worker threads); returns the values in force."""
    import cv2
    import torch

    if torch_threads is not None:
        torch.set_num_threads(int(torch_threads))
    if cv2_threads is not None:
        cv2.setNumThreads(int(cv2_threads))
    return {"torch_threads": int(torch.get_num_threads()), "cv2_threads": int(cv2.getNumThreads())}


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = build_parser().parse_args(argv)
    threads = set_host_threads(args.torch_threads, args.cv2_threads)
    if args.dry_run:
        payload = dry_run(
            version=args.version,
            policy=args.policy,
            trace_layout=args.trace_layout,
            sample=args.sample,
            embodiment=args.embodiment,
            layout_name=args.layout,
            host_reps=args.host_reps,
            profile_encode=args.profile_encode,
            dit_backend=args.dit_backend,
            mk_arena_dtype=args.mk_arena_dtype,
        )
        SO.write_perf_json(payload, Path(args.out))
        return 0
    if 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()
    from models.experimental.gr00t.tt import model as M
    from models.experimental.gr00t.tt.policy import TTPolicy

    device = M.open_model_device(
        TTPolicy(
            dtype_policy=args.policy,
            trace_layout=args.trace_layout,
            dit_backend=args.dit_backend,
            mk_arena_dtype=args.mk_arena_dtype,
        ),
        trace_region_size=args.trace_region_mb << 20,
        l1_small_size=D.DEFAULT_L1_SMALL_SIZE,
        num_command_queues=args.num_cqs,
    )
    try:
        payload = run(
            device,
            version=args.version,
            policy=args.policy,
            trace_layout=args.trace_layout,
            sample=args.sample,
            embodiment=args.embodiment,
            layout_name=args.layout,
            warmup=args.warmup,
            reps=args.reps,
            host_reps=args.host_reps,
            noise_seed=args.noise_seed,
            cache_root=Path(args.cache_root) if args.cache_root else None,
            embed_gather=args.embed_gather,
            cq1_uploads=resolve_cq1(args.cq1_uploads),
            profile_encode=args.profile_encode,
            dit_backend=args.dit_backend,
            mk_arena_dtype=args.mk_arena_dtype,
            load_ttnn_dit_weights=True if args.load_ttnn_dit_weights else None,
        )
    finally:
        D.close_gr00t_device(device)
    payload["meta"]["host_threads"] = threads
    SO.write_perf_json(payload, Path(args.out))
    return 0


__all__ = [
    "BENCH_NAME",
    "CQ1_MODES",
    "DIT_BACKENDS",
    "MK_ARENA_DTYPES",
    "EMBED_GATHERS",
    "POLICIES",
    "PRODUCTION_MEMBERS",
    "REPLAYER_MEMBERS",
    "STAGE_NAMES",
    "TRACE_LAYOUTS",
    "denoise_stage_name",
    "dry_run",
    "encode_profile",
    "golden_available",
    "host_timings",
    "import_gr00t_tt",
    "load_golden_inputs",
    "load_model",
    "make_noise",
    "resolve_cq1",
    "resolve_setup",
    "resolve_stage_replayer",
    "run",
    "time_host",
    "upload_rows",
]


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