import os
import torch
import triton
from aiter.ops.triton.sage_attention import sageattn_qk_int8_pv_fp16


PROBLEM_SIZES = [
    # "batch, num_qo_heads, qo_len, num_kv_heads, kv_len, head_dim, tensor_layout, is_causal"
    (1, 14040, 40, 14040, 40, 128, "NHD", False),
    (1, 14040, 40, 512, 40, 128, "NHD", False),
    (2, 14040, 40, 14040, 40, 128, "NHD", False),
    (2, 14040, 40, 512, 40, 128, "NHD", False),
    (2, 24, 1280, 24, 1280, 128, "HND", False),
    (1, 24, 1280, 24, 1280, 128, "HND", False),
    (1, 24, 4352, 24, 4352, 128, "HND", False),
    (2, 24, 4352, 24, 4352, 128, "HND", False),
    (1, 19440, 1, 19440, 1, 128, "NHD", False),
    (1, 2430, 40, 512, 40, 128, "NHD", False),
    (1, 57600, 1, 57600, 1, 128, "NHD", False),
    (1, 7200, 40, 512, 40, 128, "NHD", False),
    (1, 14080, 1, 14080, 1, 128, "NHD", False),
    (1, 3520, 24, 512, 24, 128, "NHD", False),
]


configs = [
    triton.testing.Benchmark(
        x_names=['batch', 'num_qo_heads', 'qo_len', 'num_kv_heads', 'kv_len', 'head_dim', 'tensor_layout', 'is_causal'],
        x_vals=PROBLEM_SIZES,
        line_arg='provider',
        line_vals=['triton'],
        line_names=['Triton'],
        styles=[('red', '-')],
        ylabel='Latency',
        xlabel='sizes',
        plot_name='QK int8 pv fp16 Performance',
        args={'dtype': torch.float16, 'device': 'cuda'},
    )
]


@triton.testing.perf_report(configs)
def bench_sageattn_qk_int8_pv_fp16(batch, num_qo_heads, qo_len, num_kv_heads, kv_len, head_dim,
                          tensor_layout, is_causal, provider, dtype=torch.float16, device='cuda'):
    print(
        f"\nBenchmark batch: {batch}, num_qo_heads: {num_qo_heads}, qo_len: {qo_len}, "
        f"num_kv_heads: {num_kv_heads}, kv_len: {kv_len}, head_dim: {head_dim}, "
        f"tensor_layout: {tensor_layout}, is_causal: {is_causal}"
    )

    q = torch.randn((batch, num_qo_heads, qo_len, head_dim), dtype=dtype, device=device)
    k = torch.randn((batch, num_kv_heads, kv_len, head_dim), dtype=dtype, device=device)
    v = torch.randn((batch, num_kv_heads, kv_len, head_dim), dtype=dtype, device=device)

    fn = lambda: sageattn_qk_int8_pv_fp16(q, k, v, tensor_layout=tensor_layout, is_causal=is_causal)
    return triton.testing.do_bench(fn)


if __name__ == "__main__":
    os.environ['AMDGPU_USE_BUFFER_OPS'] = "1"

    print(f"Triton QK Int8 PV FP16 Benchmark")
    bench_sageattn_qk_int8_pv_fp16.run(print_data=True, save_path="bench_qk_int8_pv_fp16_out")
