Accélération GPU du Game of Life : benchmarks CUDA Triton

Optimisation du Jeu de la Vie de Conway sur GPU : PyTorch, CUDA et Triton

L'automate cellulaire du Jeu de la Vie de Conway est parfaitement adapté au calcul parallèle sur GPU grâce à ses règles locales simples. Chaque cellule d'une grille N×N analyse ses 8 voisines : une cellule vivante survit avec 2 à 3 voisines vivantes, et une cellule morte naît avec exactement 3. Les tests sont effectués sur une Nvidia A40 avec une grille 216×216 (4 Go en int8). La limite théorique est de 11,5 ms par itération, déterminée par la bande passante mémoire de 696 Go/s.

Limites théoriques et calculs de base

La mise à jour d'une seule cellule nécessite le chargement de 9 octets et l'écriture de 1 octet. Avec 4 Go de données, le temps minimum est de 4 Go × 2 / 696 Go/s = 11,5 ms. La charge de calcul est minimale, la mémoire étant le goulot d'étranglement. Les bords de la grille sont ignorés pour simplifier.

PyTorch : de l'implémentation de base à torch.compile

PyTorch utilise un flou de boîte pour le comptage des voisines au lieu des convolutions standard en float32.

def gol_torch_sum(x: torch.Tensor) -> torch.Tensor:
    y = x[2:] + x[1:-1] + x[:-2]
    z = y[:, 2:] + y[:, 1:-1] + y[:, :-2]
    z = torch.nn.functional.pad(z, (1, 1, 1, 1), value=0)
    return ((x == 1) & (z == 4)) | (z == 3).to(torch.int8)

Version de base : 223 ms en raison de la surcharge des opérations individuelles. torch.compile fusionne le graphe et l'optimise : 38,1 ms (30 % du pic). Il génère automatiquement un noyau Triton.

CUDA : gestion manuelle des threads et du cache

Le noyau CUDA calcule une cellule par thread avec des blocs configurables (block_size_row × block_size_col).

__global__ void gol_kernel_i8(const int8_t* __restrict__ x_ptr,
                              int8_t* __restrict__ out_ptr,
                              int64_t rowstride, int64_t n) {
  int64_t x = blockIdx.x * blockDim.x + threadIdx.x;
  int64_t y = blockIdx.y * blockDim.y + threadIdx.y;

  if (x >= n - 2 || y >= n - 2) return;

  int8_t r00 = x_ptr[y * rowstride + x + 0 * rowstride + 0];
  // ... (8 voisines restantes)

  int8_t sum = r00 + r01 + r02 + r10 + r12 + r20 + r21 + r22;

  int8_t result = (r11 > 0) ? ((sum == 2) || (sum == 3) ? 1 : 0) : (sum == 3 ? 1 : 0);

  out_ptr[(y + 1) * rowstride + (x + 1)] = result;
}

Le cache L1 est crucial : sans lui, >55 ms. L'optimal est 1×128 (26 ms, 44 % du pic). Blocs ≤1024, multiples de 32. Les blocs carrés minimisent le périmètre, les rectangulaires exploitent le row-major.

Paramètres clés des blocs CUDA :

• Maximum 1024 threads/bloc
• Multiple de 32 pour l'occupation
• Équilibre registres et mémoire partagée
• Préférence pour 1×128 pour les tâches liées à la mémoire

Triton : vectorisation et optimisation automatique

Triton simplifie CUDA en ajoutant des opérations tensorielles. Le noyau charge des blocs 3×3 avec des masques de bord.

@triton.jit
def gol_triton_2d_kernel(x_ptr, out_ptr, row_stride: tl.int64, N: tl.int64, 
                         BLOCK_SIZE_ROW: tl.constexpr, BLOCK_SIZE_COL: tl.constexpr):
    # décalages et masques pour 3x3
    row00 = tl.load(x_ptr + row_offsets0 * row_stride + col_offsets0, 
                    mask=row_mask0 & col_mask0, other=0)
    # ... (9 chargements)
    
    sum = row00 + row01 + row02 + row10 +  row12 + row20 + row21 + row22
    result = tl.where(row11 > 0, (sum == 2) | (sum == 3), sum == 3).to(tl.int8)
    tl.store(out_ptr + row_offsets1 * row_stride + col_offsets1, result, 
             mask=row_mask1 & col_mask1)

Blocs de 1024 (8 cellules/thread, 128 threads). Triton vectorise automatiquement et utilise la mémoire partagée : 22,5 ms (51 % du pic).

Comparaison des performances

| Framework | Temps (ms) | % du Pic | Remarque |

|-----------|------------|-----------|------------|

| PyTorch | 223 | 5% | Surcharge |

| torch.compile | 38,1 | 30% | Fusion |

| CUDA | 26 | 44% | 1×128 |

| Triton | 22,5 | 51% | Vectorisé |

| Théorie | 11,5 | 100% | Mémoire |

Points clés à retenir

• Limite de 11,5 ms atteignable avec un cache parfait
• Triton en tête (22,5 ms) grâce à la vectorisation automatique
• Blocs 1×128 en CUDA optimaux pour les tâches liées à la mémoire
• torch.compile accélère PyTorch par 5,8x
• Prochaine étape : noyaux CUDA groupés avec boucles sur des groupes de cellules

— Editorial Team