gol_cpu.py

"""
Jeu de la Vie — Version CPU vectorisée (convolution 2D)

Usage :
    python gol_cpu.py [H] [W] [steps]
    python gol_cpu.py 512 512 100
"""
import torch
import torch.nn.functional as F
import time
import sys

# Kernel de Moore : 8 voisins, cellule centrale exclue
KERNEL = torch.tensor([[[[1., 1., 1.],
                          [1., 0., 1.],
                          [1., 1., 1.]]]])

def step(grid):
    """
    Une génération du Jeu de la Vie.
    grid : float32 (1, 1, H, W) avec valeurs 0.0 ou 1.0
    """
    # Padding circulaire → grille toroïdale
    padded    = F.pad(grid, (1, 1, 1, 1), mode="circular")
    neighbors = F.conv2d(padded, KERNEL)

    alive   = grid == 1.0
    survive = alive  & ((neighbors == 2) | (neighbors == 3))
    born    = ~alive & (neighbors == 3)
    return (survive | born).float()


def run(H, W, steps, seed=42):
    torch.manual_seed(seed)
    grid = (torch.rand(1, 1, H, W) < 0.3).float()

    print(f"Grille {H}×{W}  |  {steps} étapes  |  device : CPU")
    print(f"Cellules vivantes initiales : {grid.sum():.0f}")

    t0 = time.time()
    for i in range(steps):
        grid = step(grid)
    elapsed = time.time() - t0

    print(f"Cellules vivantes finales   : {grid.sum():.0f}")
    print(f"Temps total : {elapsed:.4f} s  ({elapsed/steps*1000:.3f} ms/étape)")


if __name__ == "__main__":
    H     = int(sys.argv[1]) if len(sys.argv) > 1 else 512
    W     = int(sys.argv[2]) if len(sys.argv) > 2 else 512
    steps = int(sys.argv[3]) if len(sys.argv) > 3 else 100
    run(H, W, steps)