Repository navigation
Conversation
The tile scans keep each thread's run of `items_per_thread` partials
next to each other in local memory, so the threads of a sub-group
access elements `items_per_thread` apart. With the default 8 items of
8 bytes they hit the same few banks: on an RTX 5080, the tile kernel
of a 10^7-element scan without a neutral element (`Optional{Float32}`
partials) took 11 wavefronts per local-memory access, kept the shared
pipe 93% busy and ran in 219 µs.
Pad each run of even length by one element, as CUB pads its
BlockExchange tiles, so that the runs spread over the banks: that
kernel now takes 74 µs (2.4 wavefronts per access), and the whole scan
0.57x the time, along `dims=1` 0.45x. Runs of primitive types of up to
4 bytes keep their length: LLVM accesses them as 16-byte vectors,
which conflict at most 2-way, and padding broke that alignment, which
cost an occupancy step on Metal and AMDGPU and made `Float32` scans
along `dims=1` 6-16% slower there.
The default `items_per_thread` accounts for the padding.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The tile scans (
scan_tile_kernel!,slice_tiles_kernel!) load a tile striped, then let each thread scan its own run ofitems_per_threadconsecutive elements in local memory (thread_scan!,tile_prefix!). Neighbouring threads are thereforeitems_per_threadelements apart, and with the default 8 items of 8 bytes, the threads of a warp all hit the same few banks. On an RTX 5080, the tile kernel of a 10^7-elementaccumulate!with an operator without neutral element (Optional{Float32}partials) needs 11 wavefronts per shared-memory access, keeps the shared pipe 93% busy, and takes 219 µs, against 49 µs for aFloat32cumsum moving the same bytes.This PR stores the runs at an odd stride: a run of even length gets one unused element after it, as CUB pads its BlockExchange tiles, which spreads the runs over the banks for any element size or number of banks. That kernel now takes 74 µs (2.4 wavefronts per access, shared pipe 69%, DRAM 79%), and
Float64partials go from 8.7 to 2.7 wavefronts per access. The defaultitems_per_threadaccounts for the padding.Runs of primitive types of up to 4 bytes keep their length. LLVM loads such a run as 16-byte vectors, which conflict at most 2-way, and the padding breaks their alignment: on Metal (M1) and AMDGPU (gfx1036) the scalar accesses took more registers and cost an occupancy step (
slice_tiles_kernel!from 1024 to 832 threads per threadgroup on Metal, from 16 to 12 waves on AMDGPU), makingFloat32scans alongdims=16-16% slower.findall's byte tiles were affected the same way. On CUDA those scans are bound by DRAM either way, so for these types the kernels stay as they are.Alternatives: one pad element per 32, as in CUB's
offset + (offset >> 5), assumes 32 four-byte banks and creates conflicts for odditems_per_thread; scanning the runs in registers, as CUB's ThreadScan does, took another 10% off theOptionalkernel on CUDA, but makes the 16-byte non-commutative scans of the tests fail on oneAPI and keepsitems_per_threadpartials live across the block scan.Time ratios to main, medians of 3 interleaved runs, 10^4 to 10^8 elements:
accumulate!, no neutral element,Float32dims=1accumulate!(+),Float64dims=1Float32/Int32scans,findallMetal keeps its kernels: with its 16 KiB tuning budget, 8-byte partials get 7 items per thread, whose runs are already odd. On PoCL, whose CPU has no banks, the padded scans are within the noise of a shared machine (0.94-1.08), at most a few percent slower.
A new test scans with runs of each parity and partials of 1 to 16 bytes, whole and along
dims. The accumulate and findall tests pass on CUDA, AMDGPU, PoCL, oneAPI and Metal.