Skip to content

Fit the scans' tiles in local memory for large element types - #208

Merged
maleadt merged 2 commits into
tb/reduce-subgroupsfrom
tb/scan-localmem
Nov 8, 2026
Merged

maleadt merged 2 commits into
tb/reduce-subgroupsfrom
tb/scan-localmem

Conversation

@maleadt

@maleadt maleadt commented Nov 5, 2026 •

Copy link
Copy Markdown
Member

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_thread partial results in local memory, plus one partial result per thread, and the default items_per_thread never 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 covers cumsum of ForwardDiff duals with 16 partials (136 bytes) and wider structs.

  • Large element types: an unset block size is now lowered to the largest power of two whose tile of one item per thread, plus the per-thread totals, fits local_mem_bytes. As in Fit the reductions' partial results in local memory for large element types #204 (CUB's MemBoundScaling), the size is taken from the partial-result type, which is an Optional for operators without a neutral element. The block size never goes below 2, since ScanPrefixes needs tiles of at least two elements. An explicit block_size is kept as given. This covers ScanPrefixes, SliceScan and everything built on them (accumulate!, cumsum, cumprod). The :rows slices already fit the budget, and findall's block scans work on Int32 counts.
  • One value per sub-group: when a block scan uses sub-groups, it now keeps one slot per sub-group in local memory instead of one per thread, the way Combine the values of a block with sub-group reductions #205 does for the reductions. Each thread passes its value to block_exclusive_scan directly. 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 by ptxas -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:

case before after
cumsum F32 1e6 7.6 (7.6–7.6) 7.6 (7.6–7.7) -0.8%
cumsum F32 1e8 1387.9 (1385.2–1388.9) 1388.2 (1386.4–1388.4) +0.0%
cumsum F64 1e7 265.2 (264.8–266.6) 263.7 (263.4–264.5) -0.5%
cumsum F32 1e7, no neutral 92.2 (91.2–92.3) 88.7 (87.9–89.2) -3.8%
cumsum F32 4096² dims=1 140.1 (139.7–140.2) 141.4 (140.6–141.6) +1.0%
cumsum F32 4096² dims=2 238.7 (238.5–239.4) 239.9 (238.6–240.9) +0.5%
cumsum F32 2×1e7 dims=2 274.1 (273.3–274.3) 272.3 (271.9–272.4) -0.7%
cumsum F32 1e6×8 dims=1 43.3 (43.2–43.3) 41.5 (41.5–42.2) -4.0%
cumsum F64 1000×1e4 dims=1 175.6 (175.3–176.0) 176.2 (176.0–176.9) +0.3%
findall F32 1e7 147.8 (147.4–147.9) 147.8 (147.2–148.2) -0.0%
scan 32-byte struct 1e6 54.7 (54.6–55.0) 51.9 (51.6–52.0) -5.0%
scan 96-byte struct 1e6 329.3 (328.5–329.5) 325.8 (325.7–326.4) -1.0%
scan 96-byte struct 3×1e5 dims=2 124.9 (124.7–125.5) 111.7 (110.5–112.0) -10.5%
cumsum Dual16 1e6 1716.7 (1690.3–1720.0) 1214.7 (1214.3–1214.9) -29.2%
cumsum Dual16 3×1e5 dims=2 603.4 (600.3–604.1) 472.5 (472.5–472.8) -21.7%

(The large types do compile on Blackwell before this change, because it allows about 99 KiB of shared memory.)

@maleadt
maleadt added this pull request to stack #206 November 5, 2026 11:19
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
maleadt merged commit 840f4c7 into main Nov 8, 2026
17 of 33 checks passed
@maleadt
maleadt deleted the tb/scan-localmem branch November 8, 2026 15:44
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