From a0f8018c5ed3534811b1ac89948a16096e539839 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 5 Nov 2026 10:05:08 +0100 Subject: [PATCH] Combine the values of a block with sub-group reductions The reduction kernels combined a block's values in local memory, one per thread, with a tree of log2(block size) synchronized steps. Where the backend has linear sub-groups that can shuffle the accumulator, reduce each sub-group's values with `KI.sub_group_reduce` and then the sub-groups' values, as CUB's `BLOCK_REDUCE_WARP_REDUCTIONS`, keeping one value per sub-group in local memory. This serves the whole-array, by-block, multigroup and second-pass kernels, and the ordered kernel, as the sub-group reductions combine in lane order; `ByLanes` reductions whose lanes fit a sub-group use no local memory at all. The strided kernels, whose lanes span several sub-groups, keep their trees. The tuning's new `subgroups` turns this off, as for the scans on CPU devices. On an RTX 5080 the kernel times are unchanged within noise (Float32 to ComplexF64 sums and maxima, reductions along `dims` of several shapes, in element order, and of structs of 8 to 136 bytes), so this mainly frees local memory. --- src/block.jl | 49 ++++++++++++++++++++---- src/reduce/kernels_nd.jl | 68 ++++++++++++++++----------------- src/reduce/mapreduce_1d_gpu.jl | 16 ++++---- src/reduce/mapreduce_nd.jl | 8 +++- src/reduce/mapreduce_ordered.jl | 31 ++++++++++----- src/reduce/tuning.jl | 14 +++++-- test/generic/reduce.jl | 33 ++++++++++++++++ 7 files changed, 154 insertions(+), 65 deletions(-) diff --git a/src/block.jl b/src/block.jl index a29fc42f..bda215b3 100644 --- a/src/block.jl +++ b/src/block.jl @@ -17,19 +17,54 @@ end -# The sub-group width `W` for a block scan of `T` in blocks of `block_size` threads, as a `Val`, -# or `Val(0)` for the local-memory scan. Sub-groups must be enabled by the scan tuning, linear -# (from KernelAbstractions 0.10), able to shuffle `T`, and `W ≤ block_size ≤ W²`, as one -# sub-group scans the sub-groups' totals. +# The sub-group width `W` for a block scan or reduction of `T` in blocks of `block_size` threads, +# as a `Val`, or `Val(0)` for the local-memory one. Sub-groups must be enabled by the tuning, +# linear (from KernelAbstractions 0.10), able to shuffle `T`, and `W ≤ block_size ≤ W²`, as one +# sub-group combines the sub-groups' values. @static if isdefined(KernelAbstractions, :KernelInterface) - function scan_subgroup_width(backend, ::Type{T}, block_size) where {T} - tuning(scan_tuning, backend, T).subgroups || return Val(0) + function subgroup_width(backend, ::Type{T}, block_size) where {T} KI.supports_linear_subgroups(backend) && KI.supports_shuffle(backend, T) || return Val(0) W = KI.sub_group_size(backend) return ispow2(W) && W <= block_size <= W * W ? Val(W) : Val(0) end else - scan_subgroup_width(backend, T, block_size) = Val(0) + subgroup_width(backend, T, block_size) = Val(0) +end +scan_subgroup_width(backend, ::Type{T}, block_size) where {T} = + tuning(scan_tuning, backend, T).subgroups ? subgroup_width(backend, T, block_size) : Val(0) +reduce_subgroup_width(backend, ::Type{T}, block_size) where {T} = + tuning(reduce_tuning, backend, T).subgroups ? subgroup_width(backend, T, block_size) : Val(0) + +# The local memory of a block reduction of `N` values (see `block_reduce`): a value per +# sub-group, or per thread without sub-groups +@inline block_slots(N, ::Val{W}) where {W} = W == 0 ? N : N ÷ W + +# Reduce the values `x` of the block's threads with `op` into the result of thread 0, with +# `sdata` of `block_slots(@groupsize()[1], subgroups)` elements. All threads in the block must +# call it. With a sub-group width `W` (see `reduce_subgroup_width`), each sub-group reduces its +# values, and then the sub-groups' values (as CUB's `BLOCK_REDUCE_WARP_REDUCTIONS`), in the order +# of the threads; else a tree in local memory, which needs `op` to be commutative. It does +# not synchronize after reading `sdata`, so a block that reduces again must first. +@inline function block_reduce(@context, op, sdata, x, ithread, ::Val{0}) + sdata[ithread + 0x1] = x + @synchronize() + @inline reduce_group!(@context, op, sdata, ithread) + return sdata[0x1] +end +@static if isdefined(KernelAbstractions, :KernelInterface) +@inline function block_reduce(@context, op, sdata, x, ithread, ::Val{W}) where {W} + N = @groupsize()[1] + x = KI.sub_group_reduce(op, x) + N == W && return x + nsg = N ÷ W + if ithread % W == 0x0 + sdata[ithread ÷ W + 0x1] = x + end + @synchronize() + # (every sub-group combines the sub-groups' values, as some backends (PoCL) need the + # sub-groups of a block to run the same sub-group operations; thread 0's is the result) + return KI.sub_group_reduce(op, sdata[ithread % nsg + 0x1], nsg) +end end diff --git a/src/reduce/kernels_nd.jl b/src/reduce/kernels_nd.jl index 76294548..44f7d387 100644 --- a/src/reduce/kernels_nd.jl +++ b/src/reduce/kernels_nd.jl @@ -84,17 +84,19 @@ end # `ByLanes`: each output reduces `reduce_size` consecutive elements (from the offset # of its outer index), with `lanes` threads that read them in an interleaved order; a block of -# `N` threads covers `N ÷ lanes` outputs, which are combined in local memory with log2(lanes) -# steps. +# `N` threads covers `N ÷ lanes` outputs, whose lanes are combined with sub-group operations +# where they fit a sub-group of width `W` (see `reduce_subgroup_width`), else in local memory with +# log2(lanes) steps. @kernel inbounds=true cpu=false unsafe_indices=true function by_lanes_kernel!( @Const(src), dst, f::F, op::OP, seed, init, base_offset, outer_strides, outer_sizes, - output_size, reduce_size, ::Val{lanes}, -) where {F, OP, lanes} + output_size, reduce_size, ::Val{lanes}, ::Val{W}, +) where {F, OP, lanes, W} seed = kernel_value(seed) @uniform N = @groupsize()[1] - sdata = @localmem typeof(seed) (N,) + # (no local memory where the lanes of an output are a segment of a sub-group) + sdata = @localmem typeof(seed) (lanes <= W ? 1 : N,) mapf, accop = seeded_ops(f, op, seed) iblock = @index(Group, Linear) - 0x1 @@ -115,20 +117,25 @@ end end end - sdata[ithread + 0x1] = acc - @synchronize() + total = if lanes <= W + KI.sub_group_reduce(accop, acc, lanes) + else + sdata[ithread + 0x1] = acc + @synchronize() - step = lanes ÷ 2 - while step >= 1 - if lane < step - sdata[ithread + 0x1] = accop(sdata[ithread + 0x1], sdata[ithread + step + 0x1]) + step = lanes ÷ 2 + while step >= 1 + if lane < step + sdata[ithread + 0x1] = accop(sdata[ithread + 0x1], sdata[ithread + step + 0x1]) + end + @synchronize() + step ÷= 2 end - @synchronize() - step ÷= 2 + sdata[ithread + 0x1] end if lane == 0x0 && iout < output_size - dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, sdata[ithread + 0x1]) + dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, total) end end @@ -277,11 +284,11 @@ end base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, - output_size, reduce_size, num_blocks, + output_size, reduce_size, num_blocks, subgroups::Val, ) where {F, OP} seed = kernel_value(seed) @uniform N = @groupsize()[1] - sdata = @localmem typeof(seed) (N,) + sdata = @localmem typeof(seed) (block_slots(N, subgroups),) mapf, accop = seeded_ops(f, op, seed) iblock = @index(Group, Linear) - 0x1 @@ -298,13 +305,10 @@ end j += N end - sdata[ithread + 0x1] = acc - @synchronize() - - @inline reduce_group!(@context, accop, sdata, ithread) + total = @inline block_reduce(@context, accop, sdata, acc, ithread, subgroups) if ithread == 0x0 - dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, sdata[0x1]) + dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, total) end next_iout = iout + num_blocks @@ -323,11 +327,11 @@ end base_offset, outer_strides, outer_sizes, output_size, reduce_size, num_blocks, - ::Val{W}, + ::Val{W}, subgroups::Val, ) where {F, OP, W} seed = kernel_value(seed) @uniform N = @groupsize()[1] - sdata = @localmem typeof(seed) (N,) + sdata = @localmem typeof(seed) (block_slots(N, subgroups),) mapf, accop = seeded_ops(f, op, seed) iblock = @index(Group, Linear) - 0x1 @@ -350,13 +354,10 @@ end jw += UInt32(1) end - sdata[ithread + 0x1] = acc - @synchronize() - - @inline reduce_group!(@context, accop, sdata, ithread) + total = @inline block_reduce(@context, accop, sdata, acc, ithread, subgroups) if ithread == 0x0 - dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, sdata[0x1]) + dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, total) end next_iout = iout + num_blocks @@ -377,11 +378,11 @@ end base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, - ::Val{output_size}, reduce_size, groups, + ::Val{output_size}, reduce_size, groups, subgroups::Val, ) where {F, OP, output_size} seed = kernel_value(seed) @uniform N = @groupsize()[1] - sdata = @localmem typeof(seed) (N,) + sdata = @localmem typeof(seed) (block_slots(N, subgroups),) mapf, accop = seeded_ops(f, op, seed) iblock = @index(Group, Linear) - 0x1 @@ -401,12 +402,9 @@ end j += N * groups end - sdata[ithread + 0x1] = acc - @synchronize() - - @inline reduce_group!(@context, accop, sdata, ithread) + total = @inline block_reduce(@context, accop, sdata, acc, ithread, subgroups) if ithread == 0x0 - partial[iout + igroup * output_size + 0x1] = sdata[0x1] + partial[iout + igroup * output_size + 0x1] = total end end diff --git a/src/reduce/mapreduce_1d_gpu.jl b/src/reduce/mapreduce_1d_gpu.jl index 040558d2..43a4532f 100644 --- a/src/reduce/mapreduce_1d_gpu.jl +++ b/src/reduce/mapreduce_1d_gpu.jl @@ -10,11 +10,11 @@ # `iblock` reduces the tiles `iblock`, `iblock + nblocks`, ... of `src`, and stores # `finish(op, init, dst, iblock + 1, partial)` in `dst[iblock + 1]`. @kernel inbounds=true cpu=false unsafe_indices=true function mapreduce_1d_kernel!( - @Const(src), dst, f::F, op::OP, seed, init, nblocks, ::Val{NI}, ::Val{K}, + @Const(src), dst, f::F, op::OP, seed, init, nblocks, ::Val{NI}, ::Val{K}, subgroups::Val, ) where {F, OP, NI, K} seed = kernel_value(seed) - sdata = @localmem typeof(seed) (NI,) + sdata = @localmem typeof(seed) (block_slots(NI, subgroups),) mapf, accop = seeded_ops(f, op, seed) len = length(src) @@ -42,14 +42,10 @@ end end end - sdata[ithread + 0x1] = acc - - @synchronize() - - @inline reduce_group!(@context, accop, sdata, ithread) + total = @inline block_reduce(@context, accop, sdata, acc, ithread, subgroups) if ithread == 0x0 - dst[iblock + 0x1] = finish(op, init, dst, iblock + 0x1, sdata[0x1]) + dst[iblock + 0x1] = finish(op, init, dst, iblock + 0x1, total) end end @@ -126,7 +122,9 @@ function mapreduce_1d_pass!(commutative::Bool, backend, block_size, items_per_th if commutative kernel! = mapreduce_1d_kernel!(backend, block_size) Base.inferencebarrier(kernel!)(src, dst, f, op, seed, init, blocks, Val(block_size), - Val(items_per_thread); ndrange=(block_size * blocks,)) + Val(items_per_thread), + reduce_subgroup_width(backend, typeof(seed), block_size); + ndrange=(block_size * blocks,)) else len = length(src) src, base_offset = something(buffer, (src, 0)) diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index 9734d76a..4f4c3ff2 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -135,6 +135,7 @@ function launch!(shape::ByLanes, f::F, op::OP, dst, src, buffer, backend, alg, b by_lanes_kernel!(backend, block_size)( buffer, dst, f, op, seed, init, base_offset, outer_strides, outer_sizes, dst_size, reduce_size, Val(shape.lanes), + reduce_subgroup_width(backend, typeof(seed), block_size), ndrange=(block_size * blocks,), ) end @@ -152,6 +153,7 @@ function launch!(shape::ByBlock, f::F, op::OP, dst, src, buffer, backend, alg, b Base.inferencebarrier(kernel!)( buffer, dst, f, op, seed, init, base_offset, outer_strides, outer_sizes, dst_size, reduce_size, blocks, Val(W), + reduce_subgroup_width(backend, typeof(seed), block_size), ndrange=(block_size * blocks,), ) else @@ -159,6 +161,7 @@ function launch!(shape::ByBlock, f::F, op::OP, dst, src, buffer, backend, alg, b Base.inferencebarrier(kernel!)( buffer, dst, f, op, seed, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, dst_size, reduce_size, blocks, + reduce_subgroup_width(backend, typeof(seed), block_size), ndrange=(block_size * blocks,), ) end @@ -174,6 +177,7 @@ function launch!(shape::ByBlocks, f::F, op::OP, dst, src, buffer, backend, alg, buffer, partials, f, op, seed, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, Val(dst_size), reduce_size, groups, + reduce_subgroup_width(backend, typeof(seed), block_size), ndrange=(block_size * dst_size * groups,), ) combine_groups!(backend, partials, dst, op, seed, init, dst_size, groups, false, @@ -238,7 +242,9 @@ function combine_groups!(backend, partials, dst, op::OP, seed, init, dst_size, g else by_block_kernel!(backend, block_size)( partials, dst, PartialInputs(), op, seed, init, 0, (1,), (dst_size,), (dst_size,), - (groups,), dst_size, groups, dst_size, ndrange=(block_size * dst_size,)) + (groups,), dst_size, groups, dst_size, + reduce_subgroup_width(backend, typeof(seed), block_size), + ndrange=(block_size * dst_size,)) end end diff --git a/src/reduce/mapreduce_ordered.jl b/src/reduce/mapreduce_ordered.jl index b9f39c3b..686a6aea 100644 --- a/src/reduce/mapreduce_ordered.jl +++ b/src/reduce/mapreduce_ordered.jl @@ -6,9 +6,9 @@ # Block `iblock` reduces group `g` of output `iout` (`iblock = iout + g * output_size`): the # outputs' slices are split into tiles of `N * K` elements, and each of the `groups` groups takes # a contiguous range of them (an even share, as CUB's `GridEvenShare`). In each tile, every thread -# reduces a run of `K` consecutive elements, and the block combines the runs in order (with the -# scans' block scan, which uses sub-groups where `subgroups` says so) and appends the tile's total -# to the group's. Indices and `f`, `seed` and `init` are those of the other reduction kernels (see +# reduces a run of `K` consecutive elements, and the block combines the runs in order (with +# sub-group reductions where `subgroups` says so, else the scans' block scan) and appends the +# tile's total to the group's. Indices and `f`, `seed` and `init` are those of the other reduction kernels (see # `by_block_kernel!` and `mapreduce_1d_kernel!`); with several groups, `dst` holds # `output_size × groups` partial results, for a second pass of this kernel with # `f === PartialInputs()`. With a vector width `W`, the runs of full tiles are read with `K ÷ W` @@ -24,7 +24,7 @@ ) where {F, OP, K, W} seed = kernel_value(seed) @uniform N = @groupsize()[1] - totals = @localmem typeof(seed) (N,) + totals = @localmem typeof(seed) (block_slots(N, subgroups),) mapf, accop = seeded_ops(f, op, seed) iblock = @index(Group, Linear) - 0x1 @@ -63,10 +63,8 @@ end end end - totals[ithread + 0x1] = acc - _, tile_total = block_exclusive_scan!(@context, accop, totals, seed, N, ithread, - subgroups) - total = accop(total, tile_total) + total = accop(total, ordered_block_reduce(@context, accop, totals, acc, seed, N, ithread, + subgroups)) tile += 1 end @@ -76,6 +74,21 @@ end end +# The reduction of the values `x` of the block's threads in their order, into the result of thread +# 0: with sub-groups (`block_reduce`, which is in order), or with the local-memory scan +@inline function ordered_block_reduce(@context, op, totals, x, seed, N, ithread, ::Val{0}) + totals[ithread + 0x1] = x + _, total = block_exclusive_scan!(@context, op, totals, seed, N, ithread, Val(0)) + return total +end +@inline function ordered_block_reduce(@context, op, totals, x, seed, N, ithread, + subgroups::Val) + total = @inline block_reduce(@context, op, totals, x, ithread, subgroups) + # (before the next tile's values overwrite `totals`) + @synchronize() + return total +end + # One pass of `ordered_kernel!`, of `groups` blocks per output; `dense_buffer` is whether # `src` is a dense buffer, which it can read with vector loads function ordered_pass!( @@ -90,7 +103,7 @@ function ordered_pass!( src, dst, f, op, seed, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, output_size, reduce_size, groups, Val(items_per_thread), Val(W), - scan_subgroup_width(backend, typeof(seed), block_size), + reduce_subgroup_width(backend, typeof(seed), block_size), ndrange=(block_size * output_size * groups,), ) return nothing diff --git a/src/reduce/tuning.jl b/src/reduce/tuning.jl index 118d9420..14bcb730 100644 --- a/src/reduce/tuning.jl +++ b/src/reduce/tuning.jl @@ -35,9 +35,9 @@ Values that fill unset algorithm fields for reductions on one device, as returne - `block_size`, `items_per_thread`, `switch_below`: settings for unset `BlockReduce` fields; `items_per_thread=nothing` chooses it by the size of the accumulator type, enough elements for 64 bytes, from 4 to 16. -- `local_mem_bytes`: the kernels hold a partial result per thread in local memory, so an unset - `block_size` is lowered, to a power of two, until a block's partial results fit in this many - bytes; an explicit `block_size` is kept. +- `local_mem_bytes`: an unset `block_size` is lowered, to a power of two, until a partial result + per thread fits in this many bytes, as the kernels keep without sub-groups (see `subgroups`), + and the strided ones always; an explicit `block_size` is kept. - `target_blocks`: the number of blocks a reduction aims to launch to fill the device: along `dims` when there are few outputs, and at most in the first pass of a whole-array reduction (whose second pass reduces their partial results in one block). Must be positive. @@ -56,6 +56,10 @@ Values that fill unset algorithm fields for reductions on one device, as returne flag that they are set with an integer that fills the padding after the value, up to 8 bytes (so that Compute Sanitizer's initcheck doesn't see reads of unwritten bytes there), or with a `Bool`, which some compilers handle better in local memory. The results are the same. +- `subgroups`: whether the kernels combine the values of a block with sub-group operations, + where the backend has them (from KernelAbstractions 0.10), keeping a value per sub-group + rather than per thread in local memory. CPUs, whose sub-group operations synchronize the + whole block, reduce faster without. `items_per_thread` defaults to the choice by type, and `columns_split_min_elems` to 2^18, measured on CUDA. `lanes_max_bytes` defaults to 128, beyond which AMD and Intel GPUs stop gaining @@ -73,6 +77,7 @@ Base.@kwdef struct ReduceTuning lanes_max_bytes::Int = 128 lanes_short_runs::Bool = true wide_optional_flag::Bool = true + subgroups::Bool = true end """ @@ -87,7 +92,8 @@ tuning_type(::typeof(reduce_tuning)) = ReduceTuning # with a thread per output, and long ones faster with groups of threads: with PoCL on a Ryzen # 9950X, summing an 8×10^6 matrix along `dims=1` takes half the time, and a 4096×4096 one along # `dims=1` also half the time. -cpu_reduce_tuning(; kwargs...) = ReduceTuning(; lanes_short_runs=false, lanes_max_bytes=512, kwargs...) +cpu_reduce_tuning(; kwargs...) = + ReduceTuning(; lanes_short_runs=false, lanes_max_bytes=512, subgroups=false, kwargs...) reduce_tuning(::HostBackend, ::Type) = cpu_reduce_tuning() # Elements each thread loads per tile of a whole-array reduction: enough for 64 bytes of diff --git a/test/generic/reduce.jl b/test/generic/reduce.jl index a9846843..948af251 100644 --- a/test/generic/reduce.jl +++ b/test/generic/reduce.jl @@ -1011,6 +1011,39 @@ end end +# CPU devices reduce without sub-groups (`cpu_reduce_tuning`), so on the host backend the tuning +# turns them on, to check the sub-group reductions under PoCL, which runs sub-group operations +# as barriers of the whole block: whole arrays (also in element order, over several tiles), +# runs of consecutive elements (`ByLanes`), blocks that reduce several outputs in turn +# (`ByBlock`), split reductions (`ByBlocks`), and in element order along `dims`; with and +# without a neutral element (`plus` has none) +plus(a, b) = a + b +function subgroup_reductions() + @test AK.reduce_subgroup_width(BACKEND, Int32, 256) !== Val(0) + for (sz, dims) in (((10^6,), :), ((64, 1000), 1), ((20_000, 300), 1), ((100_000, 4), 1)) + xh = rand(Int32, sz...) + x = array_from_host(xh) + for op in (+, plus) + res = AK.reduce(op, x; dims, alg=REDUCE_ALG) + @test (dims isa Colon ? res : Array(res)) == reduce(+, xh; dims) + end + end + for (sz, dims) in (((70_000,), :), ((3000, 300), 1)) + xh = [randmat2() for _ in CartesianIndices(sz)] + res = AK.reduce(matmul2, array_from_host(xh); dims, commutative=false, alg=REDUCE_ALG) + @test (dims isa Colon ? res : Array(res)) == reduce(matmul2, xh; dims, init=I2) + end +end +if HOST_KERNELS + @testset "reduce: sub-groups on the host backend" begin + @eval AK.reduce_tuning(::typeof(BACKEND), ::Type) = AK.ReduceTuning(subgroups=true) + try + Base.invokelatest(subgroup_reductions) + finally + @eval AK.reduce_tuning(::typeof(BACKEND), ::Type) = AK.cpu_reduce_tuning() + end + end +end # A reduction keeps no reference to its source, so `unsafe_free!` frees the source's memory at # once (measured where the backend's pool counts it) if @isdefined(CUDACore) && BACKEND isa CUDACore.CUDABackend && CUDACore.used_memory() !== missing