Optymalizacja automatu komórkowego "Gra w życie" na GPU: PyTorch, CUDA i Triton
Automat komórkowy Conwaya "Gra w życie" doskonale nadaje się do obliczeń równoległych na GPU dzięki prostym regułom lokalnym. Każda komórka w siatce N×N analizuje 8 sąsiadów: żywa komórka przeżywa przy 2–3 żywych sąsiadach, martwa ożywa przy dokładnie 3. Testy przeprowadzono na karcie Nvidia A40 z siatką 216×216 (4 GB w int8). Teoretyczny limit to 11,5 ms na iterację, określony przez przepustowość pamięci 696 GB/s.
Ograniczenia teoretyczne i podstawowe obliczenia
Do aktualizacji jednej komórki wymagane jest załadowanie 9 bajtów i zapis 1 bajtu. Przy 4 GB danych minimalny czas to 4 GB × 2 / 696 GB/s = 11,5 ms. Obciążenie obliczeniowe jest minimalne, wąskim gardłem jest pamięć. Pomijamy granice siatki dla uproszczenia.
PyTorch: od podstawowej implementacji do torch.compile
PyTorch używa rozmycia pudełkowego do zliczania sąsiadów zamiast standardowych splotów 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)Wersja podstawowa: 223 ms z powodu narzutu pojedynczych operacji. torch.compile łączy graf i optymalizuje: 38,1 ms (30% szczytu). Automatycznie generuje jądro Triton.
CUDA: ręczne zarządzanie wątkami i pamięcią podręczną
Jądro CUDA oblicza jedną komórkę na wątek z konfigurowalnymi blokami (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];
// ... (pozostałych 8 sąsiadów)
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;
}Pamięć podręczna L1 jest kluczowa: bez niej >55 ms. Optymalnie 1×128 (26 ms, 44% szczytu). Bloki ≤1024, wielokrotności 32. Bloki kwadratowe minimalizują obwód, prostokątne wykorzystują row-major.
Kluczowe parametry bloków CUDA:
- Maksimum 1024 wątków/blok
- Wielokrotność 32 dla zajętości
- Balans rejestrów i shared memory
- Preferencja 1×128 dla zadań ograniczonych pamięcią
Triton: wektoryzacja i automatyczna optymalizacja
Triton upraszcza CUDA, dodając operacje tensorowe. Jądro ładuje bloki 3×3 z maskami granic.
@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 i masks dla 3x3
row00 = tl.load(x_ptr + row_offsets0 * row_stride + col_offsets0,
mask=row_mask0 & col_mask0, other=0)
# ... (9 załadowań)
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)Bloki po 1024 (8 komórek/wątek, 128 wątków). Triton auto-wektoryzuje i używa shared memory: 22,5 ms (51% szczytu).
Porównanie wydajności
| Framework | Czas (ms) | % szczytu | Uwaga |
|-----------|------------|-----------|------------|
| PyTorch | 223 | 5% | Overhead |
| torch.compile | 38,1 | 30% | Fusion |
| CUDA | 26 | 44% | 1×128 |
| Triton | 22,5 | 51% | Wektoryzowane |
| Teoria | 11,5 | 100% | Pamięć |
Co jest ważne
- Limit 11,5 ms osiągalny przy idealnym buforowaniu
- Triton prowadzi (22,5 ms) dzięki auto-wektoryzacji
- Bloki 1×128 w CUDA optymalne dla ograniczeń pamięci
- torch.compile przyspiesza PyTorch 5.8x
- Dalej — grouped kernels CUDA z pętlą po grupach komórek
— Editorial Team
Brak komentarzy.