基于numba的GPU上的Floyd-Warshall算法

Question

我正在使用numba在GPU上编写优化的Floyd-Warshall算法。我需要它在10,000个矩阵的情况下在几秒钟内工作。目前，处理过程大约在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重写这种方法，并在评论中发布它;)。