Non-contiguous indexing

[ilkerkesen]

Hi Mike,

Thanks for such an effort! I think non-contiguous array indexing is the only missing feature in CuArray where we use in our models heavily. Example code snippet:

julia> a1 = rand(Float32, 4,5)
4×5 Array{Float32,2}:
 0.698076  0.795957   0.501911  0.148559  0.416837
 0.817677  0.430873   0.383991  0.443963  0.201368
 0.359252  0.210537   0.277426  0.985338  0.454552
 0.821997  0.0168654  0.453663  0.40859   0.908441

julia> c1 = CuArray(a1)
4×5 CuArray{Float32,2}:
 0.698076  0.795957   0.501911  0.148559  0.416837
 0.817677  0.430873   0.383991  0.443963  0.201368
 0.359252  0.210537   0.277426  0.985338  0.454552
 0.821997  0.0168654  0.453663  0.40859   0.908441

julia> a1[[1,20]]
2-element Array{Float32,1}:
 0.698076
 0.908441

julia> g1[[1,20]]
ERROR: MethodError: Cannot `convert` an object of type Tuple{Int64} to an object of type Array{Float64,3}
This may have arisen from a call to the constructor Array{Float64,3}(...),
since type constructors fall back to convert methods.
Stacktrace:
 [1] getindex(::GPUArrays.GPUArray{Float64,3,CUDAdrv.CuArray{Float64,3},GPUArrays.CUBackend.CUContext}, ::Array{Int64,1}) at /KUFS/scratch/ikesen16/.julia/newnode/v0.6/GPUArrays/src/abstractarray.jl:398

julia> a1[[1,2],:]
2×5 Array{Float32,2}:
 0.698076  0.795957  0.501911  0.148559  0.416837
 0.817677  0.430873  0.383991  0.443963  0.201368

julia> c1[[1,2],:]
ERROR: don't know how to handle argument of type Array{Int64,1}
Stacktrace:
 [1] cudaconvert(::Array{Int64,1}) at /KUFS/scratch/ikesen16/.julia/newnode/v0.6/CUDAnative/src/execution.jl:20
 [2] broadcast(::Function, ::Tuple{Array{Int64,1},Base.Slice{Base.OneTo{Int64}}}) at ./broadcast.jl:17
 [3] _unsafe_getindex!(::CuArray{Float32,2}, ::CuArray{Float32,2}, ::Array{Int64,1}, ::Base.Slice{Base.OneTo{Int64}}, ::Vararg{Base.Slice{Base.OneTo{Int64}},N} where N) at /KUFS/scratch/ikesen16/.julia/newnode/v0.6/CuArrays/src/indexing.jl:50
 [4] macro expansion at ./multidimensional.jl:460 [inlined]
 [5] _unsafe_getindex(::IndexLinear, ::CuArray{Float32,2}, ::Array{Int64,1}, ::Base.Slice{Base.OneTo{Int64}}) at ./multidimensional.jl:453
 [6] macro expansion at ./multidimensional.jl:442 [inlined]
 [7] _getindex at ./multidimensional.jl:438 [inlined]
 [8] getindex(::CuArray{Float32,2}, ::Array{Int64,1}, ::Colon) at ./abstractarray.jl:882

julia> a1[1:2:end,1:2:end]
2×3 Array{Float32,2}:
 0.698076  0.501911  0.416837
 0.359252  0.277426  0.454552

julia> c1[1:2:end,1:2:end]
2×3 CuArray{Float32,2}:
 0.698076  0.501911  0.416837
 0.359252  0.277426  0.454552

I think if the last case works, the other is not hard to implement. I will try to test CuArrays with my dynamic neural net benchmark examples (for now, without indexing).

[MikeInnes]

This actually works if the index is a CuArray:

julia> xs[cu[1,2], :]
2×3 CuArray{Float64,2}:
 0.213428  0.977906  0.12309 
 0.416557  0.371804  0.730723

Should be pretty easy to do that conversion automatically.

[MikeInnes]

Actually, I remember why I didn't add this earlier: Base uses scalar indexing to check bounds, so converting to CuArray will be bad for performance.

Easy to fix by doing the conversion after bounds checking; for now, ironically, cpu-array indexing will be faster, but we can fix that if cuarray indexing is needed.