Skip to content

Inline the range indexing kernel bodies and launch getindex linearly - #839

Open
unsound-mint wants to merge 2 commits into
JuliaGPU:mainfrom
unsound-mint:perf-837-range-indexing
Open

unsound-mint wants to merge 2 commits into
JuliaGPU:mainfrom
unsound-mint:perf-837-range-indexing

Conversation

@unsound-mint

@unsound-mint unsound-mint commented Oct 26, 2026 •

Copy link
Copy Markdown

Fixes #837.

Range getindex and setindex! (x[:, :, r, :], and copy(view(x, :, :, r, :)), which calls it) were 7 to 17 times slower than copyto! on the same data. The kernel bodies getindex_generated and setindex_generated were not inlined into their kernels, so every work item passed the index tuples through local memory: the PTX of getindex_kernel had a call and a __local_depot.

  • The generated functions are now @inline @generated, as permute_linearindex already is.
  • getindex launches over length(dest) instead of size(dest). The kernel only uses the linear index, so a multidimensional ndrange only adds index arithmetic (cf. KA.jl-related slowdowns #565).
  • The kernels no longer bounds-check each access: before launching, the host has already checked the indices against the indexed array and the length of the source. The setindex! source is now read relative to its firstindex. A vector source with offset indices (e.g. x[2:3, 1] = Base.IdentityUnitRange(5:6)) used to throw a BoundsError in the kernel. Without the bounds check it would have read the wrong elements.

Float32, r the second quarter of the third dimension, times in µs.

RTX 4070 Laptop GPU, first commit only, median of 300 samples with CUDA.@sync:

size(x) getindex before after setindex! before after copyto!
56×56×104×32 625 82 538 79 77
28×28×208×32 420 48 279 47 45
7×7×832×32 82 18 79 17 16

RTX 5080, median of 15 runs of 100 calls (Float16 gives the same numbers):

size(x) getindex main inline + linear + no bounds checks setindex! main inline + no bounds checks copyto!
56×56×104×32 251.0 23.9 20.8 217.6 23.9 20.9 23.9
28×28×208×32 207.8 12.5 11.0 110.2 12.6 11.0 12.6
7×7×832×32 30.6 4.1 3.7 29.9 4.1 3.6 4.2

New tests: range indexing with odd sizes (so the last workgroup is only partly used), ranges in inner and outer dimensions, strided and reversed ranges, and integer indices that drop dimensions, for getindex, view copies and setindex!; and setindex! with an offset source.

unsound-mint and others added 2 commits October 27, 2026 21:04
`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.
`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.
@maleadt
maleadt force-pushed the perf-837-range-indexing branch from 491e7f6 to fc23f30 Compare October 27, 2026 20:17

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Range getindex and setindex! are up to 17× slower than copyto!: getindex_generated/setindex_generated are not inlined

2 participants