Skip to content

Pad the runs of the scan tiles to an odd stride in local memory - #202

Open
maleadt wants to merge 1 commit into
mainfrom
tb/scan-bank-conflicts
Open

maleadt wants to merge 1 commit into
mainfrom
tb/scan-bank-conflicts

Conversation

@maleadt

@maleadt maleadt commented Oct 28, 2026

Copy link
Copy Markdown
Member

The tile scans (scan_tile_kernel!, slice_tiles_kernel!) load a tile striped, then let each thread scan its own run of items_per_thread consecutive elements in local memory (thread_scan!, tile_prefix!). Neighbouring threads are therefore items_per_thread elements 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-element accumulate! 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 a Float32 cumsum 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 Float64 partials go from 8.7 to 2.7 wavefronts per access. The default items_per_thread accounts 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), making Float32 scans along dims=1 6-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 odd items_per_thread; scanning the runs in registers, as CUB's ThreadScan does, took another 10% off the Optional kernel on CUDA, but makes the 16-byte non-commutative scans of the tests fail on oneAPI and keeps items_per_thread partials live across the block scan.

Time ratios to main, medians of 3 interleaved runs, 10^4 to 10^8 elements:

CUDA (RTX 5080) AMDGPU (gfx1036) oneAPI (A770) Metal (M1)
accumulate!, no neutral element, Float32 0.57-0.83 0.38-0.65 0.87-0.95 1.00
same along dims=1 0.45-0.96 0.34-0.64 0.90-0.94 1.00
accumulate!(+), Float64 1.00-1.02 0.74-0.95
same along dims=1 0.98-1.02 0.61-0.94
Float32 / Int32 scans, findall unchanged unchanged unchanged unchanged

Metal 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.

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

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