ex1_distributed.py

"""
Exercice 1 — f(x) = sqrt(x) + sin(x²)
Version distribuée (torch.distributed)

Usage :
    torchrun --nproc_per_node=4 ex1_distributed.py [n]
    torchrun --nproc_per_node=4 ex1_distributed.py 100000
"""
import torch
import torch.distributed as dist
import time
import sys

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

def main():
    dist.init_process_group("gloo")

    rank  = dist.get_rank()
    world = dist.get_world_size()

    n = int(sys.argv[1]) if len(sys.argv) > 1 else 100_000
    assert n % world == 0, f"n={n} doit être divisible par world_size={world}"
    n_local = n // world

    # ── Processus 0 : génère et distribue ──────────────────────────────────
    if rank == 0:
        x_global = torch.abs(torch.randn(n)) + 1e-6
        chunks   = list(torch.chunk(x_global, world))
    else:
        chunks = None

    x_local = torch.empty(n_local)
    dist.scatter(x_local, scatter_list=chunks, src=0)

    # ── Calcul local vectorisé ─────────────────────────────────────────────
    dist.barrier()
    t0 = time.time()
    result_local = f_vectorized(x_local)
    t_local = time.time() - t0
    print(f"[rank {rank}/{world}]  calcul local ({n_local} elems)  {t_local:.6f} s")

    # ── Collecte sur le processus 0 ────────────────────────────────────────
    if rank == 0:
        gather_list = [torch.empty_like(result_local) for _ in range(world)]
    else:
        gather_list = None

    dist.gather(result_local, gather_list, dst=0)

    if rank == 0:
        result_global = torch.cat(gather_list)
        print(f"\n[rank 0] Résultat global : shape={result_global.shape}")
        print(f"[rank 0] Premiers éléments : {result_global[:5]}")

    dist.destroy_process_group()

if __name__ == "__main__":
    main()