diff --git a/src/host/indexing.jl b/src/host/indexing.jl index 670cdd6b..6720e605 100644 --- a/src/host/indexing.jl +++ b/src/host/indexing.jl @@ -72,7 +72,8 @@ end Is = map(adapt(ToGPU(dest)), Is) @boundscheck checkbounds(src, Is...) - getindex_kernel(get_backend(dest))(dest, src, idims, Is...; ndrange=size(dest)) + # the kernel only uses the linear index, so don't pay for a multidimensional ndrange + getindex_kernel(get_backend(dest))(dest, src, idims, Is...; ndrange=length(dest)) return dest end @@ -86,11 +87,15 @@ end i = @index(Global, Linear) getindex_generated(dest, src, idims, i, Is...) end -@generated function getindex_generated(dest, src, idims, i, Is::Vararg{Any,N}) where {N} +# these must be inlined into their kernel, or every work item passes the index tuples +# through local memory, which is several times slower +@inline @generated function getindex_generated(dest, src, idims, i, + Is::Vararg{Any,N}) where {N} quote is = @inbounds CartesianIndices(idims)[i] @nexprs $N i -> I_i = @inbounds(Is[i][is[i]]) - val = @ncall $N getindex src i -> I_i + # the host already checked the indices against `src` + val = @inbounds @ncall $N getindex src i -> I_i @inbounds dest[i] = val end end @@ -123,12 +128,15 @@ end i = @index(Global, Linear) setindex_generated(dest, src, idims, len, i, Is...) end -@generated function setindex_generated(dest, src, idims, len, i, Is::Vararg{Any,N}) where {N} +@inline @generated function setindex_generated(dest, src, idims, len, i, + Is::Vararg{Any,N}) where {N} quote 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 + # the host already checked the indices against `dest`, and that `src` has `len` + # elements; offset by `firstindex` as a vector source can have other indices + @inbounds @ncall $N setindex! dest src[firstindex(src) + i - 1] i -> I_i return end end diff --git a/test/testsuite/indexing.jl b/test/testsuite/indexing.jl index 09103fa6..c3c38d42 100644 --- a/test/testsuite/indexing.jl +++ b/test/testsuite/indexing.jl @@ -87,6 +87,29 @@ end end end + @testset "sliced getindex and setindex! with ranges" begin + # odd sizes (so the last workgroup is only partially used), ranges in inner and outer + # dimensions, strided and reversed ranges, and integer indices that drop dimensions + # (the destination then has fewer dimensions than there are indices) + @testset "$T" for T in eltypes + x = rand(T, 7, 5, 13, 3) + @testset "$(join(Is, ", "))" for Is in ((:, :, 4:9, :), (2:6, 1:2:5, :, 2), + (:, 3, 12:-3:3, 1:2), (4, :, :, 3)) + @test compare(x -> x[Is...], AT, x) + @test compare(x -> copy(view(x, Is...)), AT, x) + y = rand(T, size(x[Is...])) + @test compare((x, y) -> (x[Is...] = y; x), AT, x, y) + end + end + end + + @testset "sliced setindex!, offset source" begin + @testset "$T" for T in eltypes + src = Base.IdentityUnitRange(5:6) # indices 5:6, not 1:2 + @test compare(x -> (x[2:3, 1] = src; x), AT, zeros(T, 4, 2)) + end + end + @allowscalar @testset "empty array" begin @testset "1D" begin Ac = zeros(Float32, 10)