import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew

def run_benchmark():
    if not torch.cuda.is_available():
        print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
        return
    device = torch.device("cuda")

    init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
    inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]

    torch_model = Model(*init_inputs).cuda()
    cuda_model = ModelNew(*init_inputs).cuda()
    torch_model.eval(); cuda_model.eval()

    print("-------------------- 精度对齐验证 --------------------")
    with torch.no_grad():
        out_torch = torch_model(*inputs)
        out_cuda = cuda_model(*inputs)
    flag = torch.allclose(out_torch, out_cuda, rtol=1e-03)
    if flag:
        print("✅ 精度对齐：两个模型的输出结果非常接近。")
    else:
        print("❌ 精度不一致！")
        print(f"最大绝对误差: {(out_torch - out_cuda).abs().max().item()}" )

    print("\n-------------------- 性能加速比测试 --------------------")
    iters = 30
    torch.cuda.synchronize(); t0 = time.time()
    for _ in range(iters):
        _ = torch_model(*inputs)
    torch.cuda.synchronize(); t_torch = (time.time() - t0) / iters

    torch.cuda.synchronize(); t0 = time.time()
    for _ in range(iters):
        _ = cuda_model(*inputs)
    torch.cuda.synchronize(); t_cuda = (time.time() - t0) / iters

    print(f"PyTorch Channel-Softmax-Affine 平均执行时间: {t_torch:.6f} 秒")
    print(f"自定义 CUDA 融合内核 平均执行时间: {t_cuda:.6f} 秒")
    sp = t_torch / t_cuda if t_cuda > 0 else 0
    if t_cuda > 0:
        print(f"加速比 (Speedup): {sp:.2f}x")
    else:
        print("CUDA 内核执行时间为0，无法计算加速比。")
    return flag, sp

if __name__ == "__main__":
    run_benchmark()
