Support WrapperGPUArray nd indexing by fusing vectorized fallback.

src/host/base.jl

@@ -301,6 +301,7 @@ end
 struct ToGPU
  array::AbstractGPUArray
 end
+ToGPU(A::WrappedArray) = ToGPU(parent(A))
 function Adapt.adapt_storage(to::ToGPU, xs::Array)
  arr = similar(to.array, eltype(xs), size(xs))
  copyto!(arr, xs)

src/host/indexing.jl

@@ -61,20 +61,24 @@ end
 
 ## vectorized indexing
 
-function vectorized_getindex(src::AbstractGPUArray, Is...)
- shape = Base.index_shape(Is...)
- dest = similar(src, shape)
+function vectorized_getindex!(dest::AbstractGPUArray, src::AbstractArray, Is...)
  any(isempty, Is) && return dest # indexing with empty array
  idims = map(length, Is)
 
  # NOTE: we are pretty liberal here supporting non-GPU indices...
- Is = map(x->adapt(ToGPU(src), x), Is)
+ Is = map(adapt(ToGPU(dest)), Is)
  @boundscheck checkbounds(src, Is...)
 
  gpu_call(getindex_kernel, dest, src, idims, Is...)
  return dest
 end
 
+function vectorized_getindex(src::AbstractGPUArray, Is...)
+ shape = Base.index_shape(Is...)
+ dest = similar(src, shape)
+ return vectorized_getindex!(dest, src, Is...)
+end
+
 @generated function getindex_kernel(ctx::AbstractKernelContext, dest, src, idims,
  Is::Vararg{Any,N}) where {N}
  quote
@@ -87,7 +91,7 @@ end
  end
 end
 
-function vectorized_setindex!(dest::AbstractGPUArray, src, Is...)
+function vectorized_setindex!(dest::AbstractArray, src, Is...)
  isempty(Is) && return dest
  idims = length.(Is)
  len = prod(idims)
@@ -101,7 +105,7 @@ function vectorized_setindex!(dest::AbstractGPUArray, src, Is...)
  end
 
  # NOTE: we are pretty liberal here supporting non-GPU indices...
- Is = map(x->adapt(ToGPU(dest), x), Is)
+ Is = map(adapt(ToGPU(dest)), Is)
  @boundscheck checkbounds(dest, Is...)
 
  gpu_call(setindex_kernel, dest, adapt(ToGPU(dest), src), idims, len, Is...;
@@ -144,6 +148,23 @@ end
  end)
 end
 
+## Vectorized index overloading for `WrappedGPUArray`
+# We'd better not to overload `getindex`/`setindex!` directly as otherwise
+# the ambiguities from the default scalar fallback become a mess.
+# The default `getindex` for `AbstractArray` follows a `similar`-`copyto!` style.
+# Thus we only dispatch the `copyto!` part (`Base._unsafe_getindex!`) to our implement.
+function Base._unsafe_getindex!(dest::AbstractGPUArray, src::AbstractArray, Is::Vararg{Union{Real, AbstractArray}, N}) where {N}
+ return vectorized_getindex!(dest, src, Base.ensure_indexable(Is)...)
+end
+# Similar for `setindex!`, its default fallback is equivalent to `copyto!`.
+# We only dispatch the `copyto!` part (`Base._unsafe_setindex!`) to our implement.
+function Base._unsafe_setindex!(::IndexStyle, A::WrappedGPUArray, x, Is::Vararg{Union{Real,AbstractArray}, N}) where N
+ return vectorized_setindex!(A, x, Base.ensure_indexable(Is)...)
+end
+# And allow one more `ReshapedArray` wrapper to handle the `_maybe_reshape` optimization.
+function Base._unsafe_setindex!(::IndexStyle, A::Base.ReshapedArray{<:Any, <:Any, <:WrappedGPUArray}, x, Is::Vararg{Union{Real,AbstractArray}, N}) where N
+ return vectorized_setindex!(A, x, Base.ensure_indexable(Is)...)
+end
 
 # find*
 

test/testsuite/indexing.jl

@@ -135,6 +135,23 @@ end
  @test compare(a->a[1:1,1:1], AT, a)
  @test compare(a->a[1:1,1:1,1:1], AT, a)
  end
+
+ @testset "getindex for WrapperGPUArray" begin
+ a = rand(Float32, 5, 5)
+ @test compare(a->a'[:, 1], AT, a)
+ @test compare(a->Base.PermutedDimsArray(a, (2, 1))[2:-1:1, 1:2], AT, a)
+ @test compare(a->LowerTriangular(a)[:], AT, a) broken=(string(AT) in ["MtlArray", "oneArray"])
+ @test compare(a->Symmetric(a, :U)[a .> 0], AT, a)
+ end
+
+ @testset "setindex! for WrapperGPUArray" for T in eltypes
+ x = AT(zeros(T, (10, 10)))'
+ y = AT(rand(T, (5, 5)))
+ x[2:6, 2:6] = y
+ @test Array(parent(x)[2:6, 2:6]) == Array(y)'
+ x[2:6, 2:6] = 1:25
+ @test Array(parent(x)[2:6, 2:6]) == reshape(1:25, 5, 5)'
+ end
 end
 
 @testsuite "indexing find" (AT, eltypes)->begin