From 5367f0d572c5c62428dc8e464ae176e2a2a3fd63 Mon Sep 17 00:00:00 2001 From: Denis Andreyev Date: Tue, 27 Oct 2026 02:17:25 +0500 Subject: [PATCH 1/2] Inline the range indexing kernel bodies and launch getindex linearly `getindex_generated` and `setindex_generated` were not inlined into their kernels, so every work item passed the index tuples through local memory. Mark both as inline. The getindex kernel only uses the linear index, so launch it over `length(dest)` instead of `size(dest)` and skip the index arithmetic of a multidimensional ndrange. On an RTX 4070 Laptop GPU, `x[:, :, r, :]` and `x[:, :, r, :] = y` on a Float32 array of size (56, 56, 104, 32) go from about 630 and 540 us to 80 us, in line with `copyto!`. Add tests for range indexing with odd sizes, strided and reversed ranges, and indices that drop dimensions. --- src/host/indexing.jl | 11 ++++++++--- test/testsuite/indexing.jl | 16 ++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/src/host/indexing.jl b/src/host/indexing.jl index 670cdd6b..144de39e 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,7 +87,10 @@ 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]]) @@ -123,7 +127,8 @@ 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] diff --git a/test/testsuite/indexing.jl b/test/testsuite/indexing.jl index 09103fa6..ac979719 100644 --- a/test/testsuite/indexing.jl +++ b/test/testsuite/indexing.jl @@ -87,6 +87,22 @@ 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 + @allowscalar @testset "empty array" begin @testset "1D" begin Ac = zeros(Float32, 10) From fc23f30db1eb0bf55463795684a1cd9e07ec4300 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 27 Oct 2026 21:04:22 +0100 Subject: [PATCH 2/2] Skip the in-kernel bounds checks of range indexing `vectorized_getindex!` and `vectorized_setindex!` check the indices against the indexed array, and the source length, before launching their kernel, so checking every access again in the kernel only costs time. Without the checks, `x[:, :, r, :]` and `x[:, :, r, :] = y` take about 12% less time on an RTX 5080, a little less than `copyto!`. The source of `setindex!` is now read relative to its `firstindex`. A vector source with offset indices used to throw a `BoundsError` in the kernel, and without the check it would read the wrong elements. --- src/host/indexing.jl | 7 +++++-- test/testsuite/indexing.jl | 7 +++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/host/indexing.jl b/src/host/indexing.jl index 144de39e..6720e605 100644 --- a/src/host/indexing.jl +++ b/src/host/indexing.jl @@ -94,7 +94,8 @@ end 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 @@ -133,7 +134,9 @@ end 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 ac979719..c3c38d42 100644 --- a/test/testsuite/indexing.jl +++ b/test/testsuite/indexing.jl @@ -103,6 +103,13 @@ 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)