Repository navigation
Conversation
The one-thread-per-slice scan along `dims` steps through its chunk one element at a time. When the slices are split into chunks, it launches only `target_blocks` blocks, and with one load in flight per thread the passes are bound by memory latency rather than bandwidth. On main, ptxas happened to unroll the loop that kept the running partial in stack slots, but not the select-based combine, which made the totals pass of an (8, 10^7) `Float32` `max` scan along `dims=2` 1.34x slower on an RTX 5080. Split slices are now stepped through eight elements at a time, loaded before they are combined and stored (as the columns reduction kernel does). That scan takes 1.82 ms instead of 3.57 ms on main and 3.98 ms before this change (RTX 5080), and 1.67 ms instead of 4.74 ms on an RTX 6000 Ada. Unsplit slices keep the old loop: they launch enough threads, and the steps made them up to 1.5x slower on the Ada.
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.
Stacked on #194.
A scan along
dimswhose slices are scanned one thread each (:columns) splits them into chunks when there are too few slices to fill the device, e.g.accumulatealongdims=2of an (8, 10^7) array. The split launches onlytarget_blocksblocks, and each thread stepped through its chunk one element at a time. Its loads don't depend on each other, but with one in flight per thread the passes were bound by memory latency, not bandwidth.On main, ptxas happened to unroll the totals pass 4x, so a few loads were in flight. That loop kept the running partial in stack slots, which #194 removes, and ptxas doesn't unroll #194's loop. That is the 1.11x of #194 on that scan (RTX 5080): its totals pass went from 1.22 ms to 1.62 ms.
This PR steps through split chunks eight elements at a time, loading them before combining and storing them (the scan may be in place), as the columns reduction kernel does. Unsplit slices keep the old loop: they launch enough threads, and the steps made them up to 1.5x slower on an RTX 6000 Ada. The (8, 10^7)
maxscan now takes 1.75 ms on the RTX 5080 (3.53 ms on main, 3.98 ms with #194 alone), and 1.66 ms instead of 4.73 ms on the RTX 6000 Ada. A test shape now splits slices into chunks that aren't a multiple of the step.Time ratios against #194,
accumulateofFloat32alongdims=2, 10^6 / 10^7 / 10^8 elements (median of 3 runs, 5 for PoCL and oneAPI):max+max+The oneAPI ratios are of
accumulate!into a preallocated array. Withaccumulate, allocating the result made the (8, n/8)maxscan of 10^7 elements vary by up to 1.11x, while the per-pass kernel times were within 1%. PoCL and Metal (where other jobs ran) are noisy: the ratios above 1 lie within the spread of the runs, and unsplit slices run the same loop as before.Tested
accumulateon CUDA, AMDGPU, PoCL, oneAPI and Metal.Separately, the tile scans (
thread_scan!,tile_prefix!) access their tiles with a stride ofitems_per_threadelements, which gives about 11-way shared-memory bank conflicts on CUDA (95% of the shared-memory pipe in Nsight Compute), so padding or a register-based thread scan, as in CUB, could speed up every tile scan.