Merge pull request #326 from JuliaGPU/tb/fix_indexing

Simplify indexing kernels and fix bug with misshape inputs.

Author: maleadt

src/host/abstractarray.jl

@@ -33,6 +33,7 @@ function Serialization.serialize(s::AbstractSerializer, t::T) where T <: Abstrac
     serialize_type(s, T)
     serialize(s, Array(t))
 end
+
 function Serialization.deserialize(s::AbstractSerializer, ::Type{T}) where T <: AbstractGPUArray
     A = deserialize(s)
     T(A)

src/host/indexing.jl

@@ -122,56 +122,55 @@ Base.@propagate_inbounds Base.setindex!(A::AbstractGPUArray, v, I::Union{Integer
 
 # generalized multidimensional indexing
 
-@generated function index_kernel(ctx::AbstractKernelContext, dest::AbstractArray, src::AbstractArray, idims, Is)
-    N = length(Is.parameters)
-    quote
-        i = @linearidx dest
-        is = CartesianIndices(idims)[i]
-        @nexprs $N i -> @inbounds I_i = Is[i][is[i]]
-        @inbounds dest[i] = @ncall $N getindex src i -> I_i
-        return
-    end
-end
-
-function Base.getindex(A::AbstractGPUArray, I...)
-    _getindex(A, to_indices(A, I)...)
-end
+Base.getindex(A::AbstractGPUArray, I...) = _getindex(A, to_indices(A, I)...)
 
 function _getindex(src::AbstractGPUArray, Is...)
     shape = Base.index_shape(Is...)
     dest = similar(src, shape)
     any(isempty, Is) && return dest # indexing with empty array
     idims = map(length, Is)
+
     AT = typeof(src).name.wrapper
     # NOTE: we are pretty liberal here supporting non-GPU indices...
-    gpu_call(index_kernel, dest, src, idims, adapt(AT, Is))
+    gpu_call(getindex_kernel, dest, src, idims, adapt(AT, Is)...)
     return dest
 end
 
-@generated function setindex_kernel!(ctx::AbstractKernelContext, dest::AbstractArray, src, idims, Is, len)
-    N = length(Is.parameters)
-    idx = ntuple(i-> :(Is[$i][is[$i]]), N)
+@generated function getindex_kernel(ctx::AbstractKernelContext, dest, src, idims,
+                                    Is::Vararg{<:Any,N}) where {N}
     quote
-        i = linear_index(ctx)
-        i > len && return
-        is = CartesianIndices(idims)[i]
-        @inbounds setindex!(dest, src[is], $(idx...))
+        i = @linearidx dest
+        is = @inbounds CartesianIndices(idims)[i]
+        @nexprs $N i -> I_i = @inbounds(Is[i][is[i]])
+        val = @ncall $N getindex src i -> I_i
+        @inbounds dest[i] = val
         return
     end
 end
 
-function Base.setindex!(A::AbstractGPUArray, v, I...)
-    _setindex!(A, v, to_indices(A, I)...)
-end
+Base.setindex!(A::AbstractGPUArray, v, I...) = _setindex!(A, v, to_indices(A, I)...)
 
 function _setindex!(dest::AbstractGPUArray, src, Is...)
     isempty(Is) && return dest
     idims = length.(Is)
     len = prod(idims)
     len==0 && return dest
+
     AT = typeof(dest).name.wrapper
     # NOTE: we are pretty liberal here supporting non-GPU sources and indices...
-    gpu_call(setindex_kernel!, dest, adapt(AT, src), idims, adapt(AT, Is), len;
+    gpu_call(setindex_kernel, dest, adapt(AT, src), idims, len, adapt(AT, Is)...;
              total_threads=len)
     return dest
 end
+
+@generated function setindex_kernel(ctx::AbstractKernelContext, dest, src, idims, len,
+                                    Is::Vararg{<:Any,N}) where {N}
+    quote
+        i = linear_index(ctx)
+        i > len && return
+        is = @inbounds CartesianIndices(idims)[i]
+        @nexprs $N i -> I_i = @inbounds(Is[i][is[i]])
+        @ncall $N setindex! dest src[i] i -> I_i
+        return
+    end
+end

test/testsuite/indexing.jl

@@ -139,4 +139,8 @@ end
         @test compare(a->a[i',:], AT, a)
         @test compare(a->a[view(i,1,:),:], AT, a)
     end
+
+    @testset "JuliaGPU/CUDA.jl#461: sliced setindex" begin
+        @test compare((X,Y)->(X[1,:] = Y), AT, zeros(2,2), ones(2))
+    end
 end