# SPDX-License-Identifier: Apache-2.0
"""mb1_eltwise — layer_norm / rms_norm / gelu / add and fused ``ttnn.linear(activation=...)`` timings at GR00T shapes.

Shapes (arch synthesis §1, perf-model-v2 §2.3): DiT activations [64, 1536] and [64, 6144]; LLM [S, 2048] and
[S, 6144] for S in {128, 256, 512, 1024}; SigLIP ViT [M, 1152] / [M, 4304] and Qwen3-VL ViT [M, 1024] / [M, 4096]
for M in {256, 352, 1792}.  Every op is timed with ``--iters`` back-to-back calls (eager) and as a trace replay.

Also measured: the per-op trace-replay floor (a [32, 32] ``ttnn.add`` — research §6 unknown 2 proxy) and the fused
``ttnn.linear(activation="gelu")`` vs ``ttnn.matmul`` + ``ttnn.gelu`` at the DiT ff1 and ViT fc1 shapes.

Usage:
    with-device.sh python models/experimental/gr00t/benchmarks/mb1_eltwise.py --out results/mb1_eltwise.json
"""

from __future__ import annotations

import argparse
import math
import os
import sys
import time
from typing import Any, Dict, List

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import mb1_common as C  # noqa: E402

NORM_SHAPES = (
    [("dit", 64, 1536)]
    + [("llm", S, 2048) for S in (128, 256, 512, 1024)]
    + [("siglip", M, 1152) for M in (256, 352, 1792)]
    + [("qwenvit", M, 1024) for M in (256, 352, 1792)]
)
ACT_SHAPES = (
    [("dit", 64, 6144)]
    + [("llm", S, 6144) for S in (128, 256, 512, 1024)]
    + [("siglip", M, 4304) for M in (256, 352, 1792)]
    + [("qwenvit", M, 4096) for M in (256, 352, 1792)]
)
FUSED_SHAPES = [
    ("dit_ff1", 64, 1536, 6144),
    ("siglip_fc1_1792", 1792, 1152, 4304),
    ("qwenvit_fc1_1792", 1792, 1024, 4096),
    ("siglip_fc1_352", 352, 1152, 4304),
]


def time_op(ttnn, device, args, row: Dict[str, Any], fn) -> Dict[str, Any]:
    t = C.measure(ttnn, device, fn, iters=args.iters, warmup=3, reps=args.reps)
    row.update(eager_us=C.r(t.eager_us), trace_us=C.r(t.trace_us), error=t.error)
    print(
        f"  {row.get('op'):>22s} {row.get('group', ''):>8s} {str(row.get('shape')):>16s} {row.get('mem', ''):>4s}: trace {row.get('trace_us')} us  eager {row.get('eager_us')} us {row.get('error') or ''}"
    )
    return row


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--out", default=os.path.join(C.default_out_dir(), "mb1_eltwise.json"))
    ap.add_argument("--iters", type=int, default=32)
    ap.add_argument("--reps", type=int, default=5)
    ap.add_argument("--quick", action="store_true")
    args = ap.parse_args()

    import torch
    import ttnn

    device = C.open_device(ttnn)
    result: Dict[str, Any] = {
        "env": C.env_report(),
        "device": C.device_report(ttnn, device),
        "args": vars(args),
        "rows": [],
    }
    rows = result["rows"]
    t0 = time.perf_counter()
    ck = C.compute_config(ttnn, "HiFi2", False)
    try:
        # ---- per-op replay floor
        a = ttnn.from_torch(
            torch.randn(1, 1, 32, 32),
            dtype=ttnn.bfloat16,
            layout=ttnn.TILE_LAYOUT,
            device=device,
            memory_config=ttnn.L1_MEMORY_CONFIG,
        )
        b = ttnn.from_torch(
            torch.randn(1, 1, 32, 32),
            dtype=ttnn.bfloat16,
            layout=ttnn.TILE_LAYOUT,
            device=device,
            memory_config=ttnn.L1_MEMORY_CONFIG,
        )
        rows.append(
            time_op(
                ttnn,
                device,
                args,
                dict(op="add_tiny_floor", group="floor", shape=[32, 32], mem="L1"),
                lambda: ttnn.add(a, b, memory_config=ttnn.L1_MEMORY_CONFIG),
            )
        )
        for iters in (8, 64, 256):
            if args.quick and iters > 64:
                continue
            tr = C.time_traced(
                ttnn, device, lambda: ttnn.add(a, b, memory_config=ttnn.L1_MEMORY_CONFIG), iters=iters, reps=args.reps
            )
            rows.append(
                dict(
                    op="add_tiny_floor",
                    group="floor",
                    shape=[32, 32],
                    mem="L1",
                    iters=iters,
                    trace_us=C.r(tr["min"]),
                    trace_us_mean=C.r(tr["mean"]),
                )
            )
            print(f"  add [32,32] x{iters} per trace: {tr['min']:.2f} us/op")
        ttnn.deallocate(a)
        ttnn.deallocate(b)

        norm_shapes = NORM_SHAPES if not args.quick else NORM_SHAPES[:3]
        act_shapes = ACT_SHAPES if not args.quick else ACT_SHAPES[:3]
        for mem_name, mem in (("L1", ttnn.L1_MEMORY_CONFIG), ("DRAM", ttnn.DRAM_MEMORY_CONFIG)):
            if args.quick and mem_name == "DRAM":
                continue
            # ---- norms
            for group, M, D in norm_shapes:
                x = ttnn.from_torch(
                    torch.randn(1, 1, M, D),
                    dtype=ttnn.bfloat16,
                    layout=ttnn.TILE_LAYOUT,
                    device=device,
                    memory_config=mem,
                )
                res = ttnn.from_torch(
                    torch.randn(1, 1, M, D),
                    dtype=ttnn.bfloat16,
                    layout=ttnn.TILE_LAYOUT,
                    device=device,
                    memory_config=mem,
                )
                gamma = ttnn.from_torch(
                    torch.randn(1, 1, 1, D), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device
                )
                beta = ttnn.from_torch(
                    torch.randn(1, 1, 1, D), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device
                )
                # ROW_MAJOR affine params must be [1, 1, D/32, 32] (layernorm_device_operation.cpp:106; tt_transformers convention)
                gamma_rm = ttnn.from_torch(
                    torch.randn(1, 1, D // 32, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=device
                )
                beta_rm = ttnn.from_torch(
                    torch.randn(1, 1, D // 32, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=device
                )
                base = dict(group=group, shape=[M, D], mem=mem_name)
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="layer_norm(w,b tile)", **base),
                        lambda: ttnn.layer_norm(
                            x, weight=gamma, bias=beta, epsilon=1e-6, memory_config=mem, compute_kernel_config=ck
                        ),
                    )
                )
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="layer_norm(w,b rm)", **base),
                        lambda: ttnn.layer_norm(
                            x, weight=gamma_rm, bias=beta_rm, epsilon=1e-6, memory_config=mem, compute_kernel_config=ck
                        ),
                    )
                )
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="layer_norm(no affine)", **base),
                        lambda: ttnn.layer_norm(x, epsilon=1e-6, memory_config=mem, compute_kernel_config=ck),
                    )
                )
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="layer_norm(+residual)", **base),
                        lambda: ttnn.layer_norm(
                            x,
                            residual_input_tensor=res,
                            weight=gamma,
                            bias=beta,
                            epsilon=1e-6,
                            memory_config=mem,
                            compute_kernel_config=ck,
                        ),
                    )
                )
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="rms_norm(w)", **base),
                        lambda: ttnn.rms_norm(
                            x, weight=gamma, epsilon=1e-6, memory_config=mem, compute_kernel_config=ck
                        ),
                    )
                )
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="rms_norm(w rm)", **base),
                        lambda: ttnn.rms_norm(
                            x, weight=gamma_rm, epsilon=1e-6, memory_config=mem, compute_kernel_config=ck
                        ),
                    )
                )
                rows.append(
                    time_op(ttnn, device, args, dict(op="add", **base), lambda: ttnn.add(x, res, memory_config=mem))
                )
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="mul_rowbcast", **base),
                        lambda: ttnn.mul(x, gamma, memory_config=mem),
                    )
                )
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="add_rowbcast", **base),
                        lambda: ttnn.add(x, beta, memory_config=mem),
                    )
                )
                for t_ in (x, res, gamma, beta, gamma_rm, beta_rm):
                    ttnn.deallocate(t_)
            # ---- activations
            for group, M, D in act_shapes:
                x = ttnn.from_torch(
                    torch.randn(1, 1, M, D),
                    dtype=ttnn.bfloat16,
                    layout=ttnn.TILE_LAYOUT,
                    device=device,
                    memory_config=mem,
                )
                y = ttnn.from_torch(
                    torch.randn(1, 1, M, D),
                    dtype=ttnn.bfloat16,
                    layout=ttnn.TILE_LAYOUT,
                    device=device,
                    memory_config=mem,
                )
                base = dict(group=group, shape=[M, D], mem=mem_name)
                rows.append(
                    time_op(
                        ttnn, device, args, dict(op="gelu(Accurate)", **base), lambda: ttnn.gelu(x, memory_config=mem)
                    )
                )
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="gelu(FastLut)", **base),
                        lambda: ttnn.gelu(x, variant=ttnn.GeluVariant.FastLut, memory_config=mem),
                    )
                )
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="gelu(Tanh)", **base),
                        lambda: ttnn.gelu(x, variant=ttnn.GeluVariant.Tanh, memory_config=mem),
                    )
                )
                rows.append(
                    time_op(ttnn, device, args, dict(op="silu", **base), lambda: ttnn.silu(x, memory_config=mem))
                )
                rows.append(
                    time_op(ttnn, device, args, dict(op="add", **base), lambda: ttnn.add(x, y, memory_config=mem))
                )
                rows.append(
                    time_op(ttnn, device, args, dict(op="mul", **base), lambda: ttnn.mul(x, y, memory_config=mem))
                )
                ttnn.deallocate(x)
                ttnn.deallocate(y)
            C.write_json(args.out, result)

        # ---- fused linear activation vs separate
        for name, M, K, N in (FUSED_SHAPES if not args.quick else FUSED_SHAPES[:1]):
            mem = ttnn.L1_MEMORY_CONFIG if M <= 64 else ttnn.DRAM_MEMORY_CONFIG
            mem_name = "L1" if M <= 64 else "DRAM"
            a = ttnn.from_torch(
                torch.randn(1, 1, M, K) * 0.5,
                dtype=ttnn.bfloat16,
                layout=ttnn.TILE_LAYOUT,
                device=device,
                memory_config=mem,
            )
            for w_dtype in ("bf16", "bfp8"):
                w = ttnn.from_torch(
                    torch.randn(1, 1, K, N) / math.sqrt(K),
                    dtype=C.ttnn_dtype(ttnn, w_dtype),
                    layout=ttnn.TILE_LAYOUT,
                    device=device,
                )
                base = dict(group=name, shape=[M, K, N], mem=mem_name, w_dtype=w_dtype)
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="matmul", **base),
                        lambda: ttnn.matmul(a, w, memory_config=mem, compute_kernel_config=ck),
                    )
                )
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="linear(act=gelu)", **base),
                        lambda: ttnn.linear(a, w, activation="gelu", memory_config=mem, compute_kernel_config=ck),
                    )
                )
                rows.append(
                    time_op(
                        ttnn,
                        device,
                        args,
                        dict(op="linear(act=silu)", **base),
                        lambda: ttnn.linear(a, w, activation="silu", memory_config=mem, compute_kernel_config=ck),
                    )
                )

                def mm_then_gelu():
                    h = ttnn.matmul(a, w, memory_config=mem, compute_kernel_config=ck)
                    g = ttnn.gelu(h, memory_config=mem)
                    ttnn.deallocate(h)
                    return g

                rows.append(time_op(ttnn, device, args, dict(op="matmul+gelu(2 ops)", **base), mm_then_gelu))
                ttnn.deallocate(w)
            ttnn.deallocate(a)
        C.write_json(args.out, result)
    finally:
        result["wall_s"] = round(time.perf_counter() - t0, 1)
        C.write_json(args.out, result)
        ttnn.close_device(device)
    C.print_table(
        rows, ["op", "group", "shape", "mem", "w_dtype", "trace_us", "eager_us", "error"], "elementwise / norm timings"
    )
    return 0


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