生命游戏的 GPU 加速：CUDA Triton 基准测试

GPU 上优化康威生命游戏：PyTorch、CUDA 与 Triton 实战

康威生命游戏元胞自动机因其简单的局部规则，非常适合并行 GPU 计算。在 N×N 网格中，每个细胞分析其 8 个邻居：活细胞在拥有 2-3 个活邻居时存活，死细胞在恰好有 3 个活邻居时复活。测试在 Nvidia A40 上进行，使用 216×216 网格（int8 格式占用 4 GB）。理论极限为每次迭代 11.5 毫秒，由 696 GB/s 的内存带宽决定。

理论极限与基础计算

更新单个细胞需要加载 9 字节并写入 1 字节。对于 4 GB 数据，最短时间为 4 GB × 2 / 696 GB/s = 11.5 毫秒。计算负载极小，内存是瓶颈。为简化起见，忽略网格边界。

PyTorch：从基础实现到 torch.compile

PyTorch 使用盒式模糊进行邻居计数，而非标准 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)

基础版本：由于单个操作的开销，耗时 223 毫秒。torch.compile 合并并优化计算图：耗时 38.1 毫秒（达到峰值的 30%）。它会自动生成 Triton 内核。

CUDA：手动线程与缓存管理

CUDA 内核每个线程计算一个细胞，可配置块大小（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 个邻居)

  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 缓存至关重要：无缓存时耗时 >55 毫秒。最优配置为 1×128（26 毫秒，44% 峰值）。块大小 ≤1024，为 32 的倍数。方形块最小化周长，矩形块利用行主序。

CUDA 块参数要点：

• 每块最多 1024 个线程
• 为提升占用率，块大小为 32 的倍数
• 平衡寄存器和共享内存使用
• 内存密集型任务首选 1×128 配置

Triton：向量化与自动优化

Triton 通过添加张量操作简化 CUDA。内核加载 3×3 块并处理边界掩码。

@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):
    # 3x3 块的偏移和掩码
    row00 = tl.load(x_ptr + row_offsets0 * row_stride + col_offsets0, 
                    mask=row_mask0 & col_mask0, other=0)
    # ... (9 次加载)
    
    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)

块大小为 1024（每个线程 8 个细胞，128 个线程）。Triton 自动向量化并使用共享内存：耗时 22.5 毫秒（51% 峰值）。

性能对比

| 框架 | 时间（毫秒） | 峰值百分比 | 备注 |

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

| PyTorch | 223 | 5% | 开销大 |

| torch.compile | 38.1 | 30% | 融合优化 |

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

| Triton | 22.5 | 51% | 向量化 |

| 理论值 | 11.5 | 100% | 内存极限 |

核心要点

• 11.5 毫秒极限 可通过完美缓存实现
• Triton 领先（22.5 毫秒），得益于自动向量化
• CUDA 中 1×128 块 对内存密集型任务最优
• torch.compile 将 PyTorch 速度提升 5.8 倍
• 下一步：使用分组 CUDA 内核循环处理细胞组

— Editorial Team