ex1_gpu.py

"""
Exercice 1 — f(x) = sqrt(x) + sin(x²)
Version GPU

Usage :
    python ex1_gpu.py [n]
    python ex1_gpu.py 100000
"""
import torch
import time
import sys

def f_cpu(x):
    return torch.sqrt(x) + torch.sin(x ** 2)

def f_gpu(x_gpu):
    return torch.sqrt(x_gpu) + torch.sin(x_gpu ** 2)

if __name__ == "__main__":
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 100_000

    if not torch.cuda.is_available():
        print("CUDA non disponible. Vérifiez votre installation PyTorch.")
        sys.exit(1)

    print(f"GPU : {torch.cuda.get_device_name(0)}")
    x = torch.abs(torch.randn(n)) + 1e-6

    # ── CPU (référence) ────────────────────────────────────────────────────
    t0 = time.time()
    res_cpu = f_cpu(x)
    t_cpu = time.time() - t0
    print(f"[CPU  n={n}]  {t_cpu:.6f} s")

    # ── GPU calcul seul (sans transfert) ──────────────────────────────────
    x_gpu = x.to("cuda")
    torch.cuda.synchronize()

    t0 = time.time()
    res_gpu = f_gpu(x_gpu)
    torch.cuda.synchronize()
    t_gpu_calc = time.time() - t0
    print(f"[GPU  n={n}  calcul seul     ]  {t_gpu_calc:.6f} s")

    # ── GPU calcul + transferts ────────────────────────────────────────────
    torch.cuda.synchronize()
    t0 = time.time()
    x_gpu2   = x.to("cuda")
    res_gpu2 = f_gpu(x_gpu2)
    torch.cuda.synchronize()
    res_cpu2 = res_gpu2.to("cpu")
    t_gpu_total = time.time() - t0
    print(f"[GPU  n={n}  calcul+transferts]  {t_gpu_total:.6f} s")

    # ── Vérification ──────────────────────────────────────────────────────
    diff = (res_cpu - res_cpu2).abs().max()
    print(f"\nMax diff CPU/GPU : {diff:.2e}  ✓")

    print(f"\nPourquoi le GPU peut être plus lent sur de petits vecteurs :")
    print("  → Overhead de lancement des kernels CUDA")
    print("  → Coût du transfert PCIe (RAM ↔ VRAM)")
    print("  → Le GPU devient rentable pour n >> 1 000 000")