Open YichengDWu opened 2 years ago
This would make a difference
using CUDA, NNlib
using Flux:gpu
using BenchmarkTools
CUDA.allowscalar(false)
a = CUDA.rand(200,30,64)
idx = rand(1:64,100000)
idx_gpu = idx |> gpu
@benchmark CUDA.@sync NNlib.gather(a, idx)
@benchmark CUDA.@sync NNlib.gather(a, idx_gpu)
But I'm not sure how much acceleration is due to no bounds checking in https://github.com/FluxML/NNlibCUDA.jl/blob/master/src/gather.jl#L52
Once we have bounds checking there, we should always automatically move the index of large size to GPU if the array is on GPU, or least suggest users to do so
To conclude: we should always put (large) index on GPU when srt
is a CUDA array.
There's no major difference as it appears to me. And moving the index to GPU could cause problems as mentioned in #411