"""Where does one training micro-step go? Arm B, micro 8 x 1024, bf16. CUDA time by operator group.

    python profile_step.py [--arm B] [--attn sdpa]
"""
import argparse, collections
import numpy as np
import torch
from torch.profiler import profile, ProfilerActivity

import train as T

ap = argparse.ArgumentParser()
ap.add_argument("--arm", default="B")
ap.add_argument("--attn", default="sdpa")
a = ap.parse_args()
counts = np.load("data_smoke/unigram_counts.npy")
m = T.Model(a.arm, counts, a.attn).cuda()
opt = torch.optim.AdamW([p for p in m.parameters() if p.requires_grad], lr=1e-4)
data = T.Data("data_smoke", 0)


def step():
    x = data.batch(8).cuda()
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = m.loss(x)
    loss.backward()
    opt.step()
    opt.zero_grad(set_to_none=True)


for _ in range(3):
    step()
torch.cuda.synchronize()
with profile(activities=[ProfilerActivity.CUDA]) as prof:
    for _ in range(3):
        step()
    torch.cuda.synchronize()
groups = collections.Counter()
for e in prof.key_averages():
    t = e.self_device_time_total if hasattr(e, "self_device_time_total") else e.self_cuda_time_total
    k = e.key.lower()
    if any(s in k for s in ("flash", "attention", "fmha", "sdpa", "efficient_attention", "softmax_backward_data")):
        g = "attention kernels"
    elif any(s in k for s in ("gemm", "matmul", "mm", "cutlass", "sm90", "sm100", "sm120", "ampere", "xmma")):
        g = "matrix multiplies (all linears incl. readout)"
    elif "log_softmax" in k or "nll" in k or "cross_entropy" in k:
        g = "cross-entropy"
    elif "adam" in k or "foreach" in k:
        g = "optimizer"
    else:
        g = "elementwise / norms / other"
    groups[g] += t
tot = sum(groups.values())
for g, t in groups.most_common():
    print(f"{g:48s} {t / 1e3 / 3:8.1f} ms/step  {t / tot * 100:5.1f}%")
print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=12))