import torch
import time

# 检查是否支持 GPU
if torch.backends.mps.is_available():
    device = torch.device("mps")
    print ("MPS device is found.")
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

# 生成大型随机矩阵
size = 4000
a_cpu = torch.randn(size, size)
b_cpu = torch.randn(size, size)

# CPU 上的矩阵相乘
start_time_cpu = time.time()
result_cpu = torch.mm(a_cpu, b_cpu)
cpu_time = time.time() - start_time_cpu

# 将矩阵移动到 GPU
a_gpu = a_cpu.to(device)
b_gpu = b_cpu.to(device)

# GPU 上的矩阵相乘
start_time_gpu = time.time()
result_gpu = torch.mm(a_gpu, b_gpu)
gpu_time = time.time() - start_time_gpu

# 将 GPU 计算结果移回 CPU
result_gpu_cpu = result_gpu.to("cpu")

# 打印计算时间
print("CPU 计算时间: {:.6f} 秒".format(cpu_time))
print("GPU 计算时间: {:.6f} 秒".format(gpu_time))

# 验证 CPU 和 GPU 计算结果是否一致
print("结果一致性:", torch.allclose(result_cpu, result_gpu_cpu))