"""Custom fused attention (tt_kernels.fused_attention) vs torch fp32 and vs the library path
(head_split + SDPA 32x192 + head_concat).  usage: attn_test.py [time] [soak N] [hifi2]"""
import sys
from pathlib import Path
from common import *
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from hamer import tt_kernels as tk
dev = open_dev()
torch.manual_seed(0)
L1 = ttnn.L1_MEMORY_CONFIG
H, D, DR, S = 16, 96, 80, 192
fid = ttnn.MathFidelity.HiFi2 if "hifi2" in sys.argv else ttnn.MathFidelity.HiFi4
q = torch.randn(1, H, S, D) * 1.5; k = torch.randn(1, H, S, D) * 1.5; v = torch.randn(1, H, S, D)
for t in (q, k, v):
    t[..., DR:] = 0
qkv_t = torch.cat([q.permute(0, 2, 1, 3).reshape(1, 1, S, H * D), k.permute(0, 2, 1, 3).reshape(1, 1, S, H * D),
                   v.permute(0, 2, 1, 3).reshape(1, 1, S, H * D)], dim=-1).to(torch.bfloat16)
qb, kb, vb = [t.to(torch.bfloat16).float() for t in (q, k, v)]
scale = DR ** -0.5
ref = torch.softmax(qb @ kb.transpose(-1, -2) * scale, -1) @ vb               # [1,H,S,D]
ref = ref.permute(0, 2, 1, 3).reshape(1, 1, S, H * D)
qkv = ttnn.from_torch(qkv_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
o = tk.fused_attention(qkv, H, D, scale, memory_config=L1, fidelity=fid)
ttnn.synchronize_device(dev)
ot = ttnn.to_torch(o).float()
print(f"custom first call OK: PCC {pcc(ot, ref):.7f} max|d| {float((ot - ref).abs().max()):.4f} mean|d| {float((ot - ref).abs().mean()):.5f}",
      "pad cols zero:", bool((ot.reshape(S, H, D)[..., DR:] == 0).all()), flush=True)
pc = ttnn.SDPAProgramConfig(compute_with_storage_grid_size=dev.compute_with_storage_grid_size(), q_chunk_size=32,
                            k_chunk_size=S, exp_approx_mode=False)
def lib():
    qq, kk, vv = tk.head_split(qkv, H, D, memory_config=L1)
    c = ttnn.transformer.scaled_dot_product_attention(qq, kk, vv, is_causal=False, scale=scale, program_config=pc, memory_config=L1)
    for t in (qq, kk, vv):
        ttnn.deallocate(t)
    cc = tk.head_concat(c, memory_config=L1); ttnn.deallocate(c)
    return cc
ol = lib(); olt = ttnn.to_torch(ol).float()
print(f"library: PCC {pcc(olt, ref):.7f} max|d| {float((olt - ref).abs().max()):.4f} mean|d| {float((olt - ref).abs().mean()):.5f}")
if "time" in sys.argv:
    print(f"custom attn {time_op(dev, lambda: tk.fused_attention(qkv, H, D, scale, memory_config=L1, fidelity=fid)):.2f} us; "
          f"library split+sdpa+concat {time_op(dev, lib):.2f} us", flush=True)
if "soak" in sys.argv:
    n = int(sys.argv[sys.argv.index("soak") + 1])
    tid = ttnn.begin_trace_capture(dev, cq_id=0)
    try:
        outs = [tk.fused_attention(qkv, H, D, scale, memory_config=L1, fidelity=fid) for _ in range(10)]
    finally:
        ttnn.end_trace_capture(dev, tid, cq_id=0)
    for i in range(n // 10):
        ttnn.execute_trace(dev, tid, cq_id=0, blocking=False)
        if i % 50 == 0:
            ttnn.synchronize_device(dev)
    ttnn.synchronize_device(dev)
    bad = sum(int(not torch.equal(ttnn.to_torch(x).float(), ot)) for x in outs)
    print(f"soak {n} OK, differing: {bad}/10")
    ttnn.release_trace(dev, tid)
ttnn.close_device(dev)