import torch
import torch.utils.benchmark as benchmark
import numpy as np
from typing import Callable

def flush_cache():
    """清除 L2 缓存。A100 的 L2 约为 40MB，这里分配 40MB 零向量并操作它。"""
    # 这里的容量根据你的 GPU 型号调整，通常 40MB-80MB 足够
    cache_size = 40 * 1024 * 1024
    x = torch.empty(cache_size, dtype=torch.int8, device='cuda')
    x.zero_()

def benchmark_operator(op_func: Callable, warmup_steps: int = 5, test_steps: int = 10):
    """
    专业的算子性能测试函数
    """
    # 1. 预热 (Warmup)
    # 排除 JIT 编译和 cuDNN autotune 的干扰
    for _ in range(warmup_steps):
        op_func()

    torch.cuda.synchronize()

    # 2. 使用 CUDA Events 进行精确计时
    start_events = [torch.cuda.Event(enable_timing=True) for _ in range(test_steps)]
    end_events = [torch.cuda.Event(enable_timing=True) for _ in range(test_steps)]

    for i in range(test_steps):
        # 3. 清除缓存 (可选，视测试需求而定)
        flush_cache()

        # 4. 插入少量时钟周期休眠以确保 GPU 队列饱和（针对极快算子）
        torch.cuda._sleep(10_000_000)

        start_events[i].record()
        op_func()
        end_events[i].record()

    # 等待所有操作完成
    torch.cuda.synchronize()

    # 计算平均时间 (单位：毫秒)
    times = [s.elapsed_time(e) for s, e in zip(start_events, end_events)]
    times_array = np.array(times)

    avg_time = np.mean(times_array)
    median_time = np.median(times_array)
    p90_time = np.percentile(times_array, 90)
    p99_time = np.percentile(times_array, 99)

    print(f"Average Execution Time: {avg_time:.4f} ms")
    print(f"Median Time: {median_time:.4f} ms")
    print(f"P90 Time: {p90_time:.4f} ms")
    print(f"P99 Time: {p99_time:.4f} ms")
    print(f"Min Time: {min(times):.4f} ms | Max Time: {max(times):.4f} ms")
    return avg_time

# --- 使用示例 ---
if __name__ == "__main__":
    # 定义你想要测试的算子，例如一个矩阵乘法
    A = torch.randn(2048, 2048, device='cuda')
    B = torch.randn(2048, 2048, device='cuda')

    def my_op():
        return torch.matmul(A, B)

    print("Benchmarking MatMul...")
    benchmark_operator(my_op)