Aceleración GPU del Game of Life: benchmarks CUDA Triton

Optimización del Juego de la Vida de Conway en GPU: PyTorch, CUDA y Triton

El autómata celular del Juego de la Vida de Conway es ideal para la computación paralela en GPU gracias a sus reglas locales simples. Cada celda en una cuadrícula de N×N analiza sus 8 vecinos: una celda viva sobrevive con 2–3 vecinos vivos, y una celda muerta revive con exactamente 3. Las pruebas se realizan en una Nvidia A40 con una cuadrícula de 216×216 (4 GB en int8). El límite teórico es de 11,5 ms por iteración, determinado por el ancho de banda de memoria de 696 GB/s.

Límites teóricos y cálculos básicos

Actualizar una sola celda requiere cargar 9 bytes y escribir 1 byte. Con 4 GB de datos, el tiempo mínimo es 4 GB × 2 / 696 GB/s = 11,5 ms. La carga computacional es mínima, siendo la memoria el cuello de botella. Los bordes de la cuadrícula se ignoran por simplicidad.

PyTorch: De la implementación básica a torch.compile

PyTorch utiliza un desenfoque de caja para contar vecinos en lugar de convoluciones estándar 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)

Versión básica: 223 ms debido a la sobrecarga de operaciones individuales. torch.compile fusiona el grafo y lo optimiza: 38,1 ms (30% del pico). Genera automáticamente un kernel de Triton.

CUDA: Gestión manual de hilos y caché

El kernel de CUDA calcula una celda por hilo con bloques 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 vecinos 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;
}

La caché L1 es crítica: sin ella, >55 ms. Lo óptimo es 1×128 (26 ms, 44% del pico). Bloques ≤1024, múltiplos de 32. Los bloques cuadrados minimizan el perímetro, los rectangulares aprovechan el orden por filas.

Parámetros clave de bloques en CUDA:

• Máximo 1024 hilos/bloque
• Múltiplo de 32 para ocupación
• Equilibrar registros y memoria compartida
• Preferencia por 1×128 para tareas limitadas por memoria

Triton: Vectorización y optimización automática

Triton simplifica CUDA añadiendo operaciones tensoriales. El kernel carga bloques 3×3 con máscaras de borde.

@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):
    # desplazamientos y máscaras para 3x3
    row00 = tl.load(x_ptr + row_offsets0 * row_stride + col_offsets0, 
                    mask=row_mask0 & col_mask0, other=0)
    # ... (9 cargas)
    
    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)

Bloques de 1024 (8 celdas/hilo, 128 hilos). Triton vectoriza automáticamente y usa memoria compartida: 22,5 ms (51% del pico).

Comparación de rendimiento

| Framework | Tiempo (ms) | % del Pico | Nota |

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

| PyTorch | 223 | 5% | Sobrecarga |

| torch.compile | 38,1 | 30% | Fusión |

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

| Triton | 22,5 | 51% | Vectorizado |

| Teoría | 11,5 | 100% | Memoria |

Conclusiones clave

• Límite de 11,5 ms alcanzable con caché perfecta
• Triton lidera (22,5 ms) gracias a la vectorización automática
• Bloques 1×128 en CUDA óptimos para tareas limitadas por memoria
• torch.compile acelera PyTorch 5,8 veces
• Próximo paso: kernels de CUDA agrupados con bucles sobre grupos de celdas

— Editorial Team