PyNNDescent on GPU

@lmcinnes

Saw your twitter response back in January about GPU implementations.

I've been implementing a cuda.jit Numba implementation of the Wasserstein distance. Ran into lots of issues with replacing functions like np.zeros, np.concatenate (anything that dynamically allocates memory), and ended up copying manual implementations (sort of for-loop C-style implementations) for e.g. np.argsort. Not so bad for Euclidean. The main issue is the (very) limited support of NumPy. I recently started learning Numba, so an expert might have a different opinion or be able to do it more easily. I was able to get Euclidean distance running on the GPU with Numba, but am still running into various issues with something more involved like Wasserstein.

For example, I haven't been able to get local array allocation working with cuda.local.array and have generally run into problems with slicing. An alternative is to pass everything in as arguments that are initialized at the top level, which is more work than just using overloaded NumPy functions that I write. I think what that means is if I want to operate on two vectors and their concatenated version, I need to pass all three in as arguments at the host (CPU) level. If I use np.diff which returns something of a different length, I also need to pass that in as an argument. The code will start to look a bit ugly, especially with nested functions - for Wasserstein I probably need about 5 additional input arguments. I'd rather be able to get it to work with cuda.local.array first and check the performance, but I might need to jump straight to this option.

Happy to share what I've tried and where I'm at if this is of interest.

No files, tests, or entry points are named. Start by reviewing the PyNNDescent implementation and the reported Numba CUDA constraints, especially local arrays, slicing, dynamic allocation, and Wasserstein distance. Define the GPU support scope and a measurable performance or correctness target before attempting implementation.