"""Quick RMS benchmark with custom configs."""
import os
import sys

import torch
from common.bench_framework import (make_bwd_benchmark_for_case,
                                    make_fwd_benchmark_for_case)
from common.diff_engine import calculate_diff

sys.path.insert(0, os.path.dirname(__file__))
from cases.rms import CASE

torch.set_default_device("cuda")

configs = [
    (512, 8, 4096),
    (1024, 8, 4096),
    (4096, 8, 4096),
    (16384, 8, 4096),
]

# Correctness check
for dim, bs, sl in configs:
    print(f"Correctness: bs={bs}, sl={sl}, D={dim}...", end=" ")
    calculate_diff(CASE, batch_size=bs, seq_len=sl, hidden_size=dim)
    print("ok")

print()

line_vals = ("naive", "naive_bw", "cuda", "cuda_bw", "speedup")
line_names = {
    "naive": "Naive (us)",
    "naive_bw": "Naive (GB/s)",
    "cuda": "CUDA (us)",
    "cuda_bw": "CUDA (GB/s)",
    "speedup": "SpeedUp (ratio)",
}

save_dir = "./results/rms_custom"
os.makedirs(save_dir, exist_ok=True)

bench = make_fwd_benchmark_for_case(
    case=CASE,
    configs=configs,
    plot_name="rms-bf16-fwd",
    dtype=torch.bfloat16,
    line_vals=line_vals,
    line_names=line_names,
)
bench.run(print_data=True, save_path=save_dir)

bench = make_bwd_benchmark_for_case(
    case=CASE,
    configs=configs,
    plot_name="rms-bf16-bwd",
    dtype=torch.bfloat16,
    line_vals=line_vals,
    line_names=line_names,
)
bench.run(print_data=True, save_path=save_dir)