All-Pairs Shortest Paths

原始题目：LeetGPU - All-Pairs Shortest Paths

题目描述

给定一个以 $N \times N$ 距离矩阵表示的有向赋权图，使用 Floyd-Warshall 算法计算所有顶点对之间的最短路径距离。矩阵以行优先顺序存储为扁平数组：$dist[i \cdot N + j]$ 是从顶点 $i$ 到顶点 $j$ 的有向边的权重。$+\infty$ 表示不存在直接边，对角线始终为零。

对于每个中间顶点 $k = 0$ 到 $N-1$（按序），更新所有对：

$$dist[i][j] = \min(dist[i][j],\ dist[i][k] + dist[k][j]) \quad \forall i, j$$

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。

约束条件

• $1 \le N \le 1{,}024$。
• 边权为非负 32 位浮点数，对角线为 0。

解题思路

Floyd-Warshall 的 $O(N^3)$ 三重循环在 GPU 上可以分块计算：将距离矩阵划分为子块，每个 block 处理一个子块。当 $N$ 较小时，$N \times N$ 的距离矩阵可以完全放入共享内存。核心递推式的数据依赖是：第 $k$ 轮更新需要第 $k$ 行和第 $k$ 列的当前值，这决定了 GPU 实现的分块和同步策略。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void floyd_step(float* dist, int N, int k) {
    int i=blockIdx.y*blockDim.y+threadIdx.y, j=blockIdx.x*blockDim.x+threadIdx.x;
    if(i<N && j<N) {
        float via = dist[i*N+k] + dist[k*N+j];
        if(via < dist[i*N+j]) dist[i*N+j] = via;
    }
}
extern "C" void solve(float* dist, int N) {
    dim3 t(16,16), b((N+15)/16,(N+15)/16);
    for(int k=0;k<N;k++){floyd_step<<<b,t>>>(dist,N,k);cudaDeviceSynchronize();}
}

Triton

import triton, triton.language as tl
@triton.jit
def floyd_step(dist_ptr, N,k, BLOCK:tl.constexpr):
    i=tl.program_id(1)*BLOCK+tl.arange(0,BLOCK); j=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK)
    mask=(i[:,None]<N)&(j[None,:]<N)
    dij=tl.load(dist_ptr+i[:,None]*N+j[None,:],mask=mask)
    dik=tl.load(dist_ptr+i[:,None]*N+k,mask=i[:,None]<N)
    dkj=tl.load(dist_ptr+k*N+j[None,:],mask=j[None,:]<N)
    tl.store(dist_ptr+i[:,None]*N+j[None,:],tl.minimum(dij,dik[:,None]+dkj[None,:]),mask=mask)