使用 numba 在 GPU 上的 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 并将其发布在评论中；）。