GPU acceleration of the Game of Life: CUDA Triton benchmarks

Optimizing Conway's Game of Life on GPU: PyTorch, CUDA, and Triton

Conway's Game of Life cellular automaton is perfectly suited for parallel GPU computing due to its simple local rules. Each cell in an N×N grid analyzes its 8 neighbors: a live cell survives with 2–3 live neighbors, and a dead cell comes to life with exactly 3. Testing is performed on an Nvidia A40 with a 216×216 grid (4 GB in int8). The theoretical limit is 11.5 ms per iteration, determined by the memory bandwidth of 696 GB/s.

Theoretical Limits and Basic Calculations

Updating a single cell requires loading 9 bytes and writing 1 byte. With 4 GB of data, the minimum time is 4 GB × 2 / 696 GB/s = 11.5 ms. The computational load is minimal, with memory being the bottleneck. Grid boundaries are ignored for simplicity.

PyTorch: From Basic Implementation to torch.compile

PyTorch uses a box blur for neighbor counting instead of standard float32 convolutions.

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)

Basic version: 223 ms due to overhead from individual operations. torch.compile merges the graph and optimizes it: 38.1 ms (30% of peak). It automatically generates a Triton kernel.

CUDA: Manual Thread and Cache Management

CUDA kernel computes one cell per thread with configurable blocks (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];
  // ... (remaining 8 neighbors)

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

L1 cache is critical: without it, >55 ms. Optimal is 1×128 (26 ms, 44% peak). Blocks ≤1024, multiples of 32. Square blocks minimize perimeter, rectangular ones leverage row-major.

Key CUDA Block Parameters:

• Maximum 1024 threads/block
• Multiple of 32 for occupancy
• Balance registers and shared memory
• Preference for 1×128 for memory-bound tasks

Triton: Vectorization and Automatic Optimization

Triton simplifies CUDA by adding tensor operations. Kernel loads 3×3 blocks with boundary masks.

@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 and masks for 3x3
    row00 = tl.load(x_ptr + row_offsets0 * row_stride + col_offsets0, 
                    mask=row_mask0 & col_mask0, other=0)
    # ... (9 loads)
    
    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)

Blocks of 1024 (8 cells/thread, 128 threads). Triton auto-vectorizes and uses shared memory: 22.5 ms (51% peak).

Performance Comparison

| Framework | Time (ms) | % of Peak | Note |

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

| PyTorch | 223 | 5% | Overhead |

| torch.compile | 38.1 | 30% | Fusion |

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

| Triton | 22.5 | 51% | Vectorized |

| Theory | 11.5 | 100% | Memory |

Key Takeaways

• 11.5 ms limit achievable with perfect caching
• Triton leads (22.5 ms) due to auto-vectorization
• 1×128 blocks in CUDA optimal for memory-bound tasks
• torch.compile speeds up PyTorch by 5.8x
• Next step: grouped CUDA kernels with loops over cell groups

— Editorial Team