#!/usr/bin/env catnip
# Sélection de device avec fallback + benchmark de latence
# Choisit CUDA si disponible (sinon CPU), puis mesure la latence moyenne
# d'un forward selon la taille de batch. Le modèle est synthétique : on
# mesure le débit du device, pas une précision.
#
# DEPS: torch

torch = import('torch')
time = import('time')

print("⇒ Device fallback + benchmark de latence")

# Fallback explicite : on préfère CUDA mais on retombe sur CPU si absent
preferred_device = 'cuda'
device = preferred_device
if preferred_device == 'cuda' and not torch.cuda.is_available() {
    print("  CUDA indisponible -> fallback CPU")
    device = 'cpu'
}
print(f"  device : {device}")

torch.manual_seed(123)
model = torch.nn.Sequential(
        torch.nn.Linear(4, 32),
        torch.nn.ReLU(),
        torch.nn.Linear(32, 3),
    ).to(device)
model.eval()
torch.set_grad_enabled(False)

# Synchronise le GPU avant de lire l'horloge : les kernels CUDA sont lancés
# de façon asynchrone, sans barrière on mesurerait le temps de lancement.
sync = () => {
    if device == 'cuda' { torch.cuda.synchronize() }
}

bench = (bs: int, runs: int): float => {
    w = 0
    while w < 10 {
        _ = model(torch.randn(tuple(bs, 4), device=device))
        w = w + 1
    }
    sync()

    start = time.perf_counter()
    i = 0
    while i < runs {
        _ = model(torch.randn(tuple(bs, 4), device=device))
        i = i + 1
    }
    sync()
    elapsed = time.perf_counter() - start
    (elapsed * 1000.0) / runs
}

batch_sizes = list(1, 16, 64, 256)
runs = 200

print()
print(f"⇒ Latence moyenne du forward ({runs} runs)")
for bs in batch_sizes {
    avg_ms = bench(bs, runs)
    print(f"  batch {bs:>4} -> {round(avg_ms, 4)} ms")
}