如何计算GPU上的成对距离矩阵

Question

The bottleneck in my code is the area where I calculate a pairwise distance matrix . Since this is the slowest part by far, I have spent much time in speeding up my code.

I have found many speedups using articles online, but the gains have been minimal. So, I am looking for a method to use my GPU to create a distance matrix in order to speed it up further. However, I know very little about using the GPU for computation. Can anyone help me do this?

In my research I have found the following, but none of them used the GPU :

1. This article was useful, but the speedups were minimal.
2. This article was informative on how to use cython and numba.

Here is an example snippet of how to calculate a pairwise distance matrix:

import numpy as np
from scipy import spatial

rows = 1000
cols = 10
mat = np.random.randn(rows, cols)
d_mat = spatial.distance.cdist(mat, mat)

My graphics card is an Nvidia Quadro M2000M

Answer 1

I was able to use this:

import numpy as np
from numba import cuda

USE_64 = True

if USE_64:
    bits = 64
    np_type = np.float64
else:
    bits = 32
    np_type = np.float32

@cuda.jit("void(float{}[:, :], float{}[:, :])".format(bits, bits))
def distance_matrix(mat, out):
    m = mat.shape[0]
    n = mat.shape[1]
    i, j = cuda.grid(2)
    d = 0
    if i < m and j < m:
        for k in range(n):
            tmp = mat[i, k] - mat[j, k]
            d += tmp * tmp
        out[i, j] = d

def gpu_dist_matrix(mat):
    rows = mat.shape[0]

    block_dim = (16, 16)
    grid_dim = (int(rows/block_dim[0] + 1), int(rows/block_dim[1] + 1))

    stream = cuda.stream()
    mat2 = cuda.to_device(np.asarray(mat, dtype=np_type), stream=stream)
    out2 = cuda.device_array((rows, rows))
    distance_matrix[grid_dim, block_dim](mat2, out2)
    out = out2.copy_to_host(stream=stream)

    return out