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)