use CuArrays indexing

Author: MikeInnes

src/construction.jl

@@ -43,10 +43,6 @@ end
 function convert(AT::Type{<: GPUArray{T1}}, A::DenseArray{T2}) where {T1, T2}
     copy!(similar(AT, T1, size(A)), T1.(A))
 end
-using Colors
-function convert(AT::Type{<: GPUArray{T1}}, A::DenseArray{T2}) where {T1 <: Colorant, T2 <: Colorant}
-    copy!(similar(AT, T1, size(A)), T1.(A))
-end
 function convert(AT::Type{<: GPUArray}, A::DenseArray{T2, N}) where {T2, N}
     copy!(similar(AT, T2, size(A)), A)
 end

src/indexing.jl

@@ -1,53 +1,64 @@
-array_convert{T, N}(t::Type{Array{T, N}}, x::Array) = convert(t, x)
-array_convert{T, N}(t::Type{Array{T, N}}, x::T) = [x]
+const _allowslow = Ref(true)
 
-function array_convert{T, N, T2}(t::Type{Array{T, N}}, x::T2)
-    arr = collect(x) # iterator
-    dims = ntuple(Val{N}) do i
-        ifelse(ndims(arr) >= i, size(arr, i), 1)
-    end
-    return reshape(map(T, arr), dims) # broadcast dims
+allowslow(flag = true) = (_allowslow[] = flag)
+
+function assertslow(op = "Operation")
+    _allowslow[] || error("$op is disabled")
+    return
 end
 
-indexlength(A, i, array::AbstractArray) = length(array)
-indexlength(A, i, array::Number) = 1
-indexlength(A, i, array::Colon) = size(A, i)
+Base.IndexStyle(::Type{<:GPUArray}) = IndexLinear()
 
-function Base.setindex!{T, N}(A::GPUArray{T, N}, value, indexes...)
-    # similarly, value should always be a julia array
-    shape = ntuple(Val{N}) do i
-        indexlength(A, i, indexes[i])
-    end
-    if !isa(value, T) # TODO, shape check errors for x[1:3] = 1
-        Base.setindex_shape_check(value, indexes...)
+function _getindex(xs::GPUArray{T}, i::Integer) where T
+    x = Array{T}(1)
+    copy!(x, 1, xs, i, 1)
+    return x[1]
+end
+
+function Base.getindex(xs::GPUArray{T}, i::Integer) where T
+    assertslow("getindex")
+    _getindex(xs, i)
+end
+
+function _setindex!(xs::GPUArray{T}, v::T, i::Integer) where T
+    x = T[v]
+    copy!(xs, i, x, 1, 1)
+    return v
+end
+
+function Base.setindex!(xs::GPUArray{T}, v::T, i::Integer) where T
+    assertslow("setindex!")
+    _setindex!(xs, v, i)
+end
+
+Base.setindex!(xs::GPUArray, v, i::Integer) = xs[i] = convert(eltype(xs), v)
+
+# Vector indexing
+
+using Base.Cartesian
+to_index(a, x) = x
+to_index(::A, x::Array{ET}) where {A, ET} = copy!(similar(A, ET, size(x)), x)
+to_index(a, x::UnitRange{<: Integer}) = convert(UnitRange{Cuint}, x)
+to_index(a, x::Base.LogicalIndex) = error("Logical indexing not implemented")
+
+@generated function index_kernel(state, dest::AbstractArray, src::AbstractArray, idims, Is)
+    N = length(Is.parameters)
+    quote
+        i = linear_index(state)
+        i > length(dest) && return
+        is = gpu_ind2sub(idims, i)
+        @nexprs $N i -> @inbounds I_i = Is[i][Int(is[i])]
+        @inbounds dest[i] = @ncall $N getindex src i -> I_i
+        return
     end
-    checkbounds(A, indexes...)
-    v = array_convert(Array{T, N}, value)
-    # since you shouldn't update GPUArrays with single indices, we simplify the interface
-    # by always mapping to ranges
-    ranges_dest = to_cartesian(A, indexes)
-    ranges_src = CartesianRange(size(v))
-
-    copy!(A, ranges_dest, v, ranges_src)
-    return
 end
 
-Base.getindex{T}(A::GPUArray{T, 0}) = Array(A)[]
-
-function Base.getindex{T, N}(A::GPUArray{T, N}, indexes...)
-    cindexes = Base.to_indices(A, indexes)
-    # similarly, value should always be a julia array
-    # We shouldn't really bother about checkbounds performance, since setindex/getindex will always be relatively slow
-    checkbounds(A, cindexes...)
-
-    shape = map(length, cindexes)
-    result = Array{T, length(shape)}(shape)
-    ranges_src = to_cartesian(A, cindexes)
-    ranges_dest = CartesianRange(shape)
-    copy!(result, ranges_dest, A, ranges_src)
-    if all(i-> isa(i, Integer), cindexes) # scalar
-        return result[]
-    else
-        return result
+
+function Base._unsafe_getindex!(dest::GPUArray, src::GPUArray, Is::Union{Real, AbstractArray}...)
+    if length(Is) == 1 && isa(first(Is), Array) && isempty(first(Is)) # indexing with empty array
+        return dest
     end
+    idims = map(length, Is)
+    gpu_call(index_kernel, dest, (dest, src, Cuint.(idims), map(x-> to_index(dest, x), Is)))
+    return dest
 end