KNN Implementation

[jenkspt]

I'm interested in a fast GPU implementation for k nearest neighbor queries -- similar to pytorch3d.ops.knn_points. Flux3D seems like an appropriate place for this to live.

Relevant pytorch3d code:
pytorch3d/csrc/knn
pytorch3d/csrc/utils/mink.cuh

Relevant Julia code:
https://github.com/JuliaParallel/rodinia/blob/master/julia_cuda/nn/nn.jl

[avik-pal]

Thanks for the references. (I can try to put together an inplementation around mid Feb). Feel free to submit a PR if you want. Even if it is not highly optimized, a simple update of the code in rodinia to latest CUDA version is acceptable to get things going.

[nirmal-suthar]

There is an internal function _nearest_neighbors which compute the nearest 1 neighbor between batched pointcloud. The GPU implementation is very naive, but the CPU version working well, if this helps.

[jenkspt]

I missed this _nearest_neighbors (gpu). I believe the advantages of the pytorch implementation are:

• Can use any distance metric (the pytorch version uses euclidean distance)
• Only the k neighbors are stored (much more memory efficient)

The rodinia findLowest apparently runs on the CPU.

I may be able to get a PR together before mid Feb. Thanks for the quick responses.

[jenkspt]

FYI: https://discourse.julialang.org/t/using-functions-in-gpu-kernel-via-kernelabstractions-jl-k-nearest-neighbor-kernel/53850

[avik-pal]

@jenkspt the link is not working

[jenkspt]

Oops. Fixed it.