GPU 上使用 numba 的 Floyd-Warshall 算法

Question

我正在使用 numba 在 GPU 上编写优化的 Floyd-Warshall 算法。如果是 10k 矩阵，我需要它在几秒钟内工作。现在处理大约在 60 秒内完成。这是我的代码：

def calcualte_distance_to_all_gpu(matrix):
    threadsperblock = (32, 32)
    blockspergrid_x = math.ceil(matrix.shape[0]/ threadsperblock[0])
    blockspergrid_y = math.ceil(matrix.shape[1] / threadsperblock[1])
    blockspergrid = (blockspergrid_x, blockspergrid_y)
    calculate_distance_to_all_cuda[blockspergrid, threadsperblock](matrix)


@cuda.jit
def calculate_distance_to_all_cuda(matrix):
    i, j = cuda.grid(2)

    N = len(matrix)
    for k in prange(N):
        if i < matrix.shape[0] and j < matrix.shape[1]:
            if matrix[i, k] + matrix[k, j] < matrix[i, j]:
                matrix[i, j] = matrix[i, k] + matrix[k, j]

老实说，我对在 GPU 上编写脚本还很陌生，所以你有什么想法可以让这段代码更快吗？我还在我的 GPU 上注意到，在处理时只有一个小峰值到 100%，然后它就不再忙了，所以问题可能出在从 CPU 向 GPU 发送数据时？如果是的话，有没有办法优化这个过程？或者也许我应该为此任务使用不同的算法？

Answer 1

事实证明，我的方法从一开始就是错误的，因为你不能直接并行化这个算法。这是一些很好的文章，如何使用代码执行此操作：

https://moorejs.github.io/APSP-in-parallel/#References

几天后，我将在评论中将此方法重写为 python numba 和 post ;)。

GPU 上使用 numba 的 Floyd-Warshall 算法

Floyd-Warshall algorithm on GPU using numba

python

optimization

gpu

floyd-warshall

numba