GPU-Beschleunigung des Game of Life: CUDA Triton-Benchmarks

Conways Game of Life auf GPU optimieren: PyTorch, CUDA und Triton

Conways Game of Life, ein zellulärer Automat, eignet sich perfekt für parallele GPU-Berechnungen dank seiner einfachen lokalen Regeln. Jede Zelle in einem N×N-Gitter analysiert ihre 8 Nachbarn: Eine lebende Zelle überlebt mit 2–3 lebenden Nachbarn, eine tote Zelle wird mit genau 3 Nachbarn lebendig. Die Tests wurden auf einer Nvidia A40 mit einem 216×216-Gitter (4 GB in int8) durchgeführt. Das theoretische Limit liegt bei 11,5 ms pro Iteration, bestimmt durch die Speicherbandbreite von 696 GB/s.

Theoretische Grenzen und Grundberechnungen

Das Aktualisieren einer einzelnen Zelle erfordert das Laden von 9 Bytes und das Schreiben von 1 Byte. Bei 4 GB Daten beträgt die Mindestzeit 4 GB × 2 / 696 GB/s = 11,5 ms. Die Rechenlast ist minimal, der Speicher ist der Engpass. Gittergrenzen werden der Einfachheit halber ignoriert.

PyTorch: Von der Basisimplementierung zu torch.compile

PyTorch verwendet einen Box-Blur für die Nachbarzählung anstelle von Standard-Float32-Faltungen.

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)

Basisversion: 223 ms aufgrund von Overhead durch einzelne Operationen. torch.compile fusioniert den Graphen und optimiert ihn: 38,1 ms (30 % des Maximums). Es generiert automatisch einen Triton-Kernel.

CUDA: Manuelles Thread- und Cache-Management

Der CUDA-Kernel berechnet eine Zelle pro Thread mit konfigurierbaren Blöcken (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];
  // ... (verbleibende 8 Nachbarn)

  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;
}

Der L1-Cache ist entscheidend: Ohne ihn >55 ms. Optimal ist 1×128 (26 ms, 44 % Maximum). Blöcke ≤1024, Vielfache von 32. Quadratische Blöcke minimieren den Umfang, rechteckige nutzen Row-Major aus.

Wichtige CUDA-Block-Parameter:

• Maximal 1024 Threads/Block
• Vielfaches von 32 für Auslastung
• Register und Shared Memory ausbalancieren
• Bevorzugung von 1×128 für speichergebundene Aufgaben

Triton: Vektorisierung und automatische Optimierung

Triton vereinfacht CUDA durch Tensoroperationen. Der Kernel lädt 3×3-Blöcke mit Grenzmasken.

@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):
    # Offsets und Masken für 3x3
    row00 = tl.load(x_ptr + row_offsets0 * row_stride + col_offsets0, 
                    mask=row_mask0 & col_mask0, other=0)
    # ... (9 Ladevorgänge)
    
    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)

Blöcke von 1024 (8 Zellen/Thread, 128 Threads). Triton vektorisiert automatisch und nutzt Shared Memory: 22,5 ms (51 % Maximum).

Leistungsvergleich

| Framework | Zeit (ms) | % des Maximums | Hinweis |

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

| PyTorch | 223 | 5% | Overhead |

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

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

| Triton | 22,5 | 51% | Vektorisiert |

| Theorie | 11,5 | 100% | Speicher |

Wichtige Erkenntnisse

• 11,5 ms Limit mit perfektem Caching erreichbar
• Triton führt (22,5 ms) dank Auto-Vektorisierung
• 1×128 Blöcke in CUDA optimal für speichergebundene Aufgaben
• torch.compile beschleunigt PyTorch um das 5,8-fache
• Nächster Schritt: Gruppierte CUDA-Kernel mit Schleifen über Zellengruppen

— Editorial Team