From b4f53a6da9f37de40245c9f1e9211fa39a856d83 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 28 Oct 2026 14:23:21 +0100 Subject: [PATCH] Keep several loads in flight per thread in the columns scan kernel 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. --- src/accumulate/accumulate_nd.jl | 37 +++++++++++++++++++++++++++++---- test/generic/accumulate.jl | 5 +++-- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/src/accumulate/accumulate_nd.jl b/src/accumulate/accumulate_nd.jl index 7dd7d446..ebe5588e 100644 --- a/src/accumulate/accumulate_nd.jl +++ b/src/accumulate/accumulate_nd.jl @@ -214,6 +214,24 @@ end end end +# The elements a thread of the columns kernel loads at a time +const SLICE_COLUMNS_ITEMS = 8 + +# The elements `load(src[i + j * stride])` for `j` in `0:N-1` +@inline load_step(src, load, seed, i, stride, ::Val{N}) where {N} = + ntuple(j -> lift(seed, load(@inbounds src[i + (j - 1) * stride])), Val(N)) + +# Combine the elements `xs`, at positions `k` onwards of a slice, into `running` in turn, storing +# each prefix in `v` unless it is `nothing` (recursing over the tuple keeps it in registers, where +# a loop indexing it can leave it in a stack slot) +@inline scan_step!(op, v, index, k, running, inclusive, ::Tuple{}) = running +@inline function scan_step!(op, v, index, k, running, inclusive, xs::Tuple) + isnothing(v) || inclusive || store_partial!(v, index(k), running) + running = op(running, first(xs)) + isnothing(v) || !inclusive || store_partial!(v, index(k), running) + return scan_step!(op, v, index, k + 1, running, inclusive, Base.tail(xs)) +end + # One thread per chunk of a slice; thread `t` takes chunk `c` of slice `(i, o)` with # `t = i + n_inner * (c + nchunks * o)`, the chunk's position in the array of totals @kernel inbounds=true cpu=false unsafe_indices=true function slice_columns_kernel!( @@ -234,12 +252,23 @@ end k1 = min(k0 + chunk_len, len) running = isnothing(carries) ? start : carries[tid + 1] + index = k -> base + k * n_inner + 1 + # (a thread's loads don't depend on each other, but they are only issued together in + # unrolled steps, which load before they store as the scan may be in place. Split slices + # take them: a split launches too few threads to hide the latency of one load at a time. + # Unsplit slices have enough threads, and there the steps are slower on an RTX 6000 Ada) + k = k0 + while nchunks > 1 && k + (SLICE_COLUMNS_ITEMS - 1) < k1 + xs = load_step(src, load, seed, index(k), n_inner, Val(SLICE_COLUMNS_ITEMS)) + running = scan_step!(op, isnothing(totals) ? v : nothing, index, k, running, + inclusive, xs) + k += SLICE_COLUMNS_ITEMS + end if isnothing(totals) - serial_scan!(op, v, src, load, k -> base + k * n_inner + 1, k0:k1 - 1, running, seed, - inclusive) + serial_scan!(op, v, src, load, index, k:k1 - 1, running, seed, inclusive) else - for k in k0:k1 - 1 - running = op(running, lift(seed, load(src[base + k * n_inner + 1]))) + for j in k:k1 - 1 + running = op(running, lift(seed, load(src[index(j)]))) end totals[tid + 1] = running end diff --git a/test/generic/accumulate.jl b/test/generic/accumulate.jl index 615cca03..b475f03c 100644 --- a/test/generic/accumulate.jl +++ b/test/generic/accumulate.jl @@ -211,7 +211,8 @@ end # short slices, or many neighbouring ones) and a block per slice scanning tiles of it, each # with the slices whole or (for the largest arrays) split into chunks whose totals are scanned # in turn. Tiles of exactly one and two blocks' worth, and one element more, for both block - # sizes. + # sizes. Slices split into chunks of a length that isn't a multiple of the columns kernel's + # step. for neutral in (I2, nothing), inclusive in (true, false), block_size in (64, 256) S = neutral === nothing ? typeof(AK.Optional{Mat2}()) : Mat2 tile = !TEST_KERNELS ? 2 * block_size : @@ -222,7 +223,7 @@ end ((3, tile), 2), ((3, tile + 1), 2), ((3, 2tile), 2), ((2tile + 1, 3), 1), ((32, 20_000), 2), ((64, 3, 500), 2), ((9, 4, 3000), 3), ((3, 70_000), 2), ((8, 6), 2), ((1, 5000, 2), 2), ((3, 200_000), 2), ((600_000, 1), 1), - ((40, 300), 1)) + ((40, 300), 1), ((32, 12_345), 2)) xh = [randmat2() for _ in CartesianIndices(sz)] init = randmat2() y = array_from_host(xh)