Skip to content

Speed up findall with CartesianIndex items - #140

Draft
maleadt wants to merge 2 commits into
mainfrom
tb/findall-cartesian
Draft

maleadt wants to merge 2 commits into
mainfrom
tb/findall-cartesian

Conversation

@maleadt

@maleadt maleadt commented Oct 1, 2026

Copy link
Copy Markdown
Member

Part of #135.

findall on a matrix returns CartesianIndexes by default (items=keys(A)). On CUDA that took a fixed 13-19 µs more than returning Ints (items=LinearIndices(A)), even for tiny masks. At 4096×4096 it took 440 µs against 185 µs.

The cause was the conversion from a linear position to a CartesianIndex. Base does it with checked, signed 64-bit integer divisions. GPUs have no hardware instruction for those and emulate them with long instruction sequences. In the scatter pass of findall, each thread converts the positions of its run one after another, so this cost was paid several times per thread.

This PR converts with unsigned divisions instead: 32-bit ones when the index space has fewer than 2^32 elements, 64-bit ones otherwise. It also drops the division-by-zero check. That check can never fire here, because a position only exists when no axis is empty. Offset and stepped axes are handled as before.

RTX 5080, mask with half its elements true. Minimum of 50 synchronized runs (CUDA.@elapsed), variants interleaved, best of 3 rounds, on a GPU shared with at most one other process. µs:

mask keys, main keys, this PR LinearIndices
100×10 38.2 24.2 23.2
1000×100 42.8 25.4 24.0
4096×4096 438.2 268.2 185.3

The small masks now cost the same as with Int items. At 4096×4096 the remaining difference comes from writing 16-byte indices instead of 8-byte ones.

On an Apple M1 and an Intel Iris Xe, the timings stay the same within noise. On the M1, Cartesian findall costs as much as gathering precomputed CartesianIndexes from an array, so the divisions were never its bottleneck there.

There are no API changes. Tested on CUDA, OpenCL (PoCL), the CPU back-end, Metal, oneAPI and AMDGPU. Direct tests of the conversion cover offset, descending and 0- and 1-dimensional indices, and index spaces just below and above 2^32 elements.

@maleadt
maleadt marked this pull request as draft October 1, 2026 15:40
@maleadt

maleadt commented Oct 1, 2026

Copy link
Copy Markdown
Member Author

Not going to merge this one yet because it may prove unnecessary once we have GPUCompiler-level integer narrowing (TBD).

`findall` with `CartesianIndex` items (`keys(A)`, the default for
matrices) converted each selected position with Base's
`CartesianIndices` indexing, which divides with checked signed 64-bit
divisions. GPUs emulate those in long instruction sequences, and in the
scatter pass every thread converts the positions of its run one after
the other, so Cartesian `findall` took a fixed 13-19 µs more than with
`Int` items, even on small masks.

Divide unsigned, in 32 bits where the length fits, and without the zero
check: no axis is empty where a position exists.

`findall` of a mask, half true, minimum of 50 synchronized runs
(`CUDA.@elapsed`), RTX 5080, µs:

| mask      | `keys` before | after | `LinearIndices` |
|-----------|---------------|-------|-----------------|
| 100×10    | 38.2          | 24.2  | 23.2            |
| 1000×100  | 42.8          | 25.4  | 24.0            |
| 4096×4096 | 438.2         | 268.2 | 185.3           |

The remaining difference at 4096×4096 is the larger output: 16-byte
indices instead of 8-byte ones.
CartesianIndices of about 2^32 elements have lengths and positions that overflow a 32-bit `Int`,
so the reference `ci[p]` in the test threw a BoundsError on x86.
@maleadt
maleadt force-pushed the tb/findall-cartesian branch from 9887576 to 96e1e34 Compare October 23, 2026 13:45

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.

1 participant