Repository navigation
Fit the scans' tiles in local memory for large element types - #208
Merged
Merged
Conversation
maleadt
force-pushed
the
tb/scan-localmem
branch
from
November 5, 2026 12:04
c3bdfcb to
07cc971
Compare
The scan kernels hold a tile of partial results and a partial result per thread in local memory, and an unset `items_per_thread` only goes down to 1, so in blocks of 256 threads elements of more than 64 bytes exceeded the 32 KiB budget, and those of more than 96 bytes failed to compile on CUDA devices before Blackwell (48 KiB of static shared memory). ForwardDiff duals with 16 partials are 136 bytes. As for the reductions, lower an unset block size to the largest power of two whose tile of one item per thread and per-thread totals fit the tuning's `local_mem_bytes`, sized by the partial-result type; but not below 2, for `ScanPrefixes`' tiles of two elements. This covers `ScanPrefixes`, `SliceScan` (whose block sizes are capped by the algorithm's) and so `accumulate!`, `cumsum` and `cumprod`; the `:rows` slices already fit the budget, and `findall`'s block scans are of `Int32` counts. An explicit `block_size` is kept. x-ref JuliaGPU/CUDA.jl#3516
With sub-groups, the block scan of the scan kernels scans each sub-group's values with `KI.sub_group_scan` and then the sub-groups' totals, but the kernels still stored every thread's value in local memory for it to read back, a slot per thread. Pass the values to the block scan instead, which keeps one slot per sub-group (`block_slots`), as the block reductions do; the local-memory scan stores them itself. This serves the tile scans and reductions of `ScanPrefixes`, `SliceScan`'s tiles and `findall`'s compaction, and drops two barriers: one waited for every thread to read its own slot, and the last one only served callers that scan again, which synchronize themselves. The tuning still sizes tiles next to a value per thread, which the backends without sub-groups, and the CPUs, need. As for the reductions, the host backend's tests turn sub-groups on to check them under PoCL.
maleadt
force-pushed
the
tb/scan-localmem
branch
from
November 6, 2026 09:59
07cc971 to
73cdf3f
Compare
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 scans have the same problem #204 fixes for the reductions (JuliaGPU/CUDA.jl#3516). Each scan kernel keeps a tile of
block_size * items_per_threadpartial results in local memory, plus one partial result per thread, and the defaultitems_per_threadnever goes below 1. With the default 256 threads, any element over 64 bytes overruns the 32 KiB budget, and any element over 96 bytes fails to compile on CUDA devices before Blackwell (48 KiB of static shared memory). That coverscumsumof ForwardDiff duals with 16 partials (136 bytes) and wider structs.local_mem_bytes. As in Fit the reductions' partial results in local memory for large element types #204 (CUB'sMemBoundScaling), the size is taken from the partial-result type, which is anOptionalfor operators without a neutral element. The block size never goes below 2, sinceScanPrefixesneeds tiles of at least two elements. An explicitblock_sizeis kept as given. This coversScanPrefixes,SliceScanand everything built on them (accumulate!,cumsum,cumprod). The:rowsslices already fit the budget, andfindall's block scans work onInt32counts.block_exclusive_scandirectly. This removes two barriers, and for large elements it shrinks the scan kernels' shared memory further (on sm_75, the largest scan kernel went from 32 KiB to 20 KiB). The tuning still sizes tiles as if there were one value per thread, because backends without sub-groups, and CPUs, need that. The host backend's tests turn sub-groups on, so this path is also checked under PoCL.On sm_75, before this change, 38 of the 83 scan kernels compiled for structs of 64 to 2080 bytes and 16-partial duals (whole arrays and along
dims, with and without a neutral element) were rejected byptxas -arch=sm_75. After it, all 117 assemble.Kernel time on an RTX 5080, in µs. Each value is the device time of one call, best of 20; the table shows the median of 3 interleaved runs per side, with the range in parentheses:
(The large types do compile on Blackwell before this change, because it allows about 99 KiB of shared memory.)