Finish whole-array reductions on the device

src/reduce/mapreduce_1d_gpu.jl

@@ -1,12 +1,20 @@
+# Whole-array reduction on the GPU, in at most two launches (as CUB's `DeviceReduce`): a first
+# pass of at most `max_blocks` blocks (the tuning's `target_blocks`), each reducing tiles of
+# `block_size * items_per_thread` elements in turn (a grid-stride loop) to one partial result,
+# and a second pass of one block that reduces those partials. Each pass stores its values with
+# `_finish`, so the last one can write the result into a device array.
+
 # NI and K are compile-time values so the local-memory size is static and the load loop unrolls.
 # `neutral` seeds each thread's partial result (an empty `_Lane` when `op` has no known neutral
-# element); `f === _Partials()` when `src` holds partial results of an earlier pass.
+# element); `f === _Partials()` when `src` holds partial results of an earlier pass. Block
+# `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_block!(
-    @Const(src), dst, f, op, neutral, ::Val{NI}, ::Val{K},
+    @Const(src), dst, f, op, neutral, init, nblocks, ::Val{NI}, ::Val{K},
 ) where {NI, K}
 
     sdata = @localmem typeof(neutral) (NI,)
-    f, op = _lanefuncs(f, op, neutral)
+    f_lanes, op_lanes = _lanefuncs(f, op, neutral)
     len = length(src)
 
     # NOTE: for many index calculations in this library, computation using zero-indexing leads to
@@ -18,101 +26,103 @@
     iblock = @index(Group, Linear) - 0x1
     ithread = @index(Local, Linear) - 0x1
 
-    # Consecutive threads load consecutive elements, while each thread advances by NI.
+    # Consecutive threads load consecutive elements, while each thread advances by NI; full
+    # tiles need no bounds checks, and only the last tile can be partial.
     acc = neutral
-    for s in 0x0:(K - 0x1)
-        idx = iblock * (NI * K) + s * NI + ithread
-        if idx < len
-            acc = op(acc, f(src[idx + 0x1]))
+    tile = iblock * (NI * K)
+    while tile + NI * K <= len
+        for s in 0x0:(K - 0x1)
+            acc = op_lanes(acc, f_lanes(src[tile + s * NI + ithread + 0x1]))
+        end
+        tile += nblocks * (NI * K)
+    end
+    if tile < len
+        for s in 0x0:(K - 0x1)
+            idx = tile + s * NI + ithread
+            if idx < len
+                acc = op_lanes(acc, f_lanes(src[idx + 0x1]))
+            end
         end
     end
     sdata[ithread + 0x1] = acc
 
     @synchronize()
 
-    @inline reduce_group!(@context, op, sdata, ithread)
+    @inline reduce_group!(@context, op_lanes, sdata, ithread)
 
     if ithread == 0x0
-        dst[iblock + 0x1] = sdata[0x1]
+        dst[iblock + 0x1] = _finish(op, init, dst, iblock + 0x1, sdata[0x1])
     end
 end
 
+# The number of blocks of the first pass over `len` elements
+_mapreduce_1d_blocks(len, block_size, items_per_thread, max_blocks) =
+    min(cld(len, block_size * items_per_thread), max_blocks)
 
-# Reduce the non-empty `src` to a host value; `neutral` is the partial-result seed of
-# `_reduce_seed`, `init` is a value or `_NoInit()`, and `partials` holds `_mapreduce_1d_partials`
-# partial results (`nothing` if none are needed).
+# Reduce the non-empty `src`; `neutral` is the partial-result seed of `_reduce_seed`, `init` is a
+# value, `_NoInit()` or (with `dst`) `_Fold()`, and `partials` holds `_mapreduce_1d_partials`
+# partial results (`nothing` if none are needed). Without `dst`, returns the result as a host
+# value; with `dst`, stores it in `dst[1]` on the device.
 function mapreduce_1d_gpu(
     f::F, op::OP, src::MapReduceSource, backend::Backend;
     init,
     neutral,
     block_size::Int,
     items_per_thread::Int,
+    max_blocks::Int,
     partials::Union{Nothing, AbstractArray},
     switch_below::Int,
+    dst::Union{Nothing, AbstractArray}=nothing,
 ) where {F, OP}
     @argcheck 1 <= block_size <= 1024
     @argcheck ispow2(block_size)
     @argcheck items_per_thread >= 1
+    @argcheck max_blocks >= 1
     @argcheck switch_below >= 0
 
-    P = typeof(neutral)
     f_host, op_host = _lanefuncs(f, op, neutral)
-
-    # Degenerate cases
     len = length(src)
-    # `f` may index device arrays too
-    len == 1 && return _finish(op, init, nothing, 0,
-                               @allowscalar(op_host(neutral, f_host(src[1]))))
-    if len < switch_below
-        h_src = _host_copy(src)
-        return _finish(op, init, nothing, 0, Base.mapreduce(f_host, op_host, h_src; init=neutral))
-    end
-
-    # Each block handles `items_per_thread * block_size` elements.
-    num_per_block = items_per_thread * block_size
-    blocks = (len + num_per_block - 1) ÷ num_per_block
 
-    dst = partials
-
-    # Later the kernel will be compiled for views anyways, so use same types for arrays.
-    src_view = _mapreduce_1d_src_view(src)
-    dst_view = @view dst[1:blocks]
-
-    kernel! = _mapreduce_block!(backend, block_size)
-    kernel!(src_view, dst_view, f, op, neutral, Val(block_size), Val(items_per_thread);
-            ndrange=(block_size * blocks,))
-
-    # As long as we still have blocks to process, swap between the src and dst pointers at
-    # the beginning of the first and second halves of dst
-    len = blocks
-    if len < switch_below
-        h_src = Vector(@view(dst[1:len]))
-        return _finish(op, init, nothing, 0, Base.reduce(op_host, h_src; init=neutral))
+    # Degenerate cases, finished on the host
+    if dst === nothing
+        # `f` may index device arrays too
+        len == 1 && return _finish(op, init, nothing, 0,
+                                   @allowscalar(op_host(neutral, f_host(src[1]))))
+        if len < switch_below
+            h_src = _host_copy(src)
+            return _finish(op, init, nothing, 0,
+                           Base.mapreduce(f_host, op_host, h_src; init=neutral))
+        end
     end
 
-    # Now all src elements have been passed through f; just do final reduction, no map needed
-    p1 = @view dst[1:len]
-    p2 = @view dst[blocks + 1:end]
-
-    while len > 1
-        blocks = (len + num_per_block - 1) ÷ num_per_block
-
-        # Each block produces one reduced value
-        kernel!(p1, p2, _Partials(), op, neutral, Val(block_size), Val(items_per_thread);
+    kernel! = _mapreduce_block!(backend, block_size)
+    NI, K = Val(block_size), Val(items_per_thread)
+    # (later passes see the source as a view, so use one for the first pass too)
+    src_view = _mapreduce_1d_src_view(src)
+    blocks = _mapreduce_1d_blocks(len, block_size, items_per_thread, max_blocks)
+
+    # The last pass stores the result: in `dst` with `init` applied, or the bare partial in the
+    # slot after the first pass's partials, for the host
+    last_dst, last_init = dst === nothing ?
+        (@view(partials[blocks + 1:blocks + 1]), _NoFinish()) : (dst, init)
+    if blocks == 1
+        kernel!(src_view, last_dst, f, op, neutral, last_init, 1, NI, K;
+                ndrange=(block_size,))
+    else
+        p = @view partials[1:blocks]
+        kernel!(src_view, p, f, op, neutral, _NoFinish(), blocks, NI, K;
                 ndrange=(block_size * blocks,))
-        len = blocks
-
-        if len < switch_below
-            h_src = Vector(@view(p2[1:len]))
-            return _finish(op, init, nothing, 0, Base.reduce(op_host, h_src; init=neutral))
+        if dst === nothing && blocks < switch_below
+            return _finish(op, init, nothing, 0, Base.reduce(op_host, Vector(p); init=neutral))
         end
-
-        p1, p2 = p2, p1
-        p1 = @view p1[1:len]
+        # (the second pass reads the partials in unrolled steps of up to 16 loads per thread: in
+        # one step for the default settings)
+        K2 = Val(clamp(cld(max_blocks, block_size), 1, 16))
+        kernel!(p, last_dst, _Partials(), op, neutral, last_init, 1, NI, K2;
+                ndrange=(block_size,))
     end
-
-    # The GPU kernel reduced all elements to one; apply init
-    return _finish(op, init, nothing, 0, @allowscalar(p1[1]))
+    dst === nothing || return dst
+    return _finish(op, init, nothing, 0, @allowscalar(last_dst[1]))
 end
 
 _host_copy(src::AbstractArray) = Array(src)
@@ -129,9 +139,16 @@ _on_host(x) = x
 _mapreduce_1d_src_view(src::AbstractArray) = @view src[1:end]
 _mapreduce_1d_src_view(src::Base.Broadcast.Broadcasted) = src
 
-# The number of partial results of `mapreduce_1d_gpu` over `len` elements (two buffers of one per
-# block of the first pass), or 0 when it needs none
-function _mapreduce_1d_partials(len, block_size, items_per_thread, switch_below)
+# The number of partial results of `mapreduce_1d_gpu` over `len` elements: one per block of the
+# first pass, and for a host result (`device=false`) a slot for the last partial; 0 when it needs
+# none
+function _mapreduce_1d_partials(len, block_size, items_per_thread, max_blocks, switch_below;
+                                device::Bool=false)
+    if device
+        len <= 1 && return 0
+        blocks = _mapreduce_1d_blocks(len, block_size, items_per_thread, max_blocks)
+        return blocks == 1 ? 0 : blocks
+    end
     (len <= 1 || len < switch_below) && return 0
-    return 2 * cld(len, block_size * items_per_thread)
+    return _mapreduce_1d_blocks(len, block_size, items_per_thread, max_blocks) + 1
 end

src/reduce/mapreduce_nd.jl

@@ -185,14 +185,33 @@ end
 # The outputs per block of the columns kernel
 _columns_cols(dst_size) = min(nextpow(2, dst_size), COLUMNS_COLS)
 
+# The source of a reduction along `dims` into a single output as a linearly indexed collection
+# of its elements, for the whole-array kernels (`mapreduce_1d_gpu`): a `Broadcasted` object, or a
+# contiguous range of a dense buffer; `nothing` for other sources.
+_mapreduce_nd_linear(src::Base.Broadcast.Broadcasted, dims_valid) = src
+function _mapreduce_nd_linear(src::AbstractArray, dims_valid)
+    layout = _mapreduce_strided_layout(src)
+    isnothing(layout) && return nothing
+    buffer, base_offset, src_strides = layout
+    (rstrides, _, nr), _ = _mapreduce_segments(size(src), src_strides, dims_valid)
+    nr == 1 && rstrides[1] == 1 || return nothing
+    return base_offset == 0 && buffer === src ? src :
+           view(buffer, base_offset + 1:base_offset + length(src))
+end
+
 # The scratch of a reduction along `dims` into `dst_size` outputs: the partial results of a
-# `:multigroup` reduction or of a `:columns` one in two passes, with the partial-result type of the
-# seed `neutral` (see `_reduce_seed`)
+# `:multigroup` reduction, of a `:columns` one in two passes or of the whole-array kernels, with
+# the partial-result type of the seed `neutral` (see `_reduce_seed`)
 function _mapreduce_nd_sizes(src, backend, alg, ::Type{A}, neutral, dims_valid, dst_size) where {A}
     alg isa BlockReduce && A !== Union{} || return (;)
     src_sizes = size(src)
     len = Base.prod(src_sizes[d] for d in dims_valid; init=1)
     (len <= 1 || dst_size == 0) && return (;)
+    if dst_size == 1 && _mapreduce_nd_linear(src, dims_valid) !== nothing
+        n = _mapreduce_1d_partials(len, alg.block_size, alg.items_per_thread,
+                                   reduce_tuning(backend, A).target_blocks, 0; device=true)
+        return n == 0 ? (;) : (; partials=_buffer(typeof(neutral), n))
+    end
     layout = _mapreduce_strided_layout(src)
     isnothing(layout) && return (;)
     segs = _mapreduce_segments(src_sizes, layout[3], dims_valid)
@@ -274,6 +293,17 @@ function mapreduce_nd!(
         return dst
     end
 
+    # A single output of a contiguous source: the whole-array kernels, finishing on the device
+    if dst_size == 1
+        linear = _mapreduce_nd_linear(src, dims_valid)
+        if linear !== nothing
+            mapreduce_1d_gpu(f, op, linear, backend; init, neutral, block_size,
+                             items_per_thread=alg.items_per_thread, max_blocks=target_blocks,
+                             partials=get(bufs, :partials, nothing), switch_below=0, dst)
+            return dst
+        end
+    end
+
     # The stride-based fast paths below index a flat buffer at
     # `buffer[base_offset + Σ coordᵈ·strideᵈ + 1]`. This works for any source backed
     # by a single dense column-major buffer (dense arrays, but also strided views,

src/reduce/reduce.jl

@@ -178,7 +178,7 @@ function _mapreduce_setup(f::F, op::OP, src, backend, init, neutral, acctype, di
         # (the buffer's name does not depend on the length, so that results infer)
         sizes = if a isa BlockReduce && seed !== nothing
             n = _mapreduce_1d_partials(length(src), a.block_size, a.items_per_thread,
-                                       a.switch_below)
+                                       reduce_tuning(backend, A).target_blocks, a.switch_below)
             (; partials=_buffer(typeof(seed), n))
         else
             (;)
@@ -363,16 +363,17 @@ function _mapreduce_whole(f::F, op::OP, src, backend, alg, ::Type{A}; init, neut
     _check_acctype(op, f, A)
 
     # The result, `op(init, partial)` or `partial`, has the accumulator type
-    return convert(A, _mapreduce_whole_run(f, op, src, backend, alg; init, neutral, partials))
+    return convert(A, _mapreduce_whole_run(f, op, src, backend, alg, A; init, neutral, partials))
 end
 
-function _mapreduce_whole_run(f::F, op::OP, src, backend, alg; init, neutral,
-                              partials) where {F, OP}
+function _mapreduce_whole_run(f::F, op::OP, src, backend, alg, ::Type{A}; init, neutral,
+                              partials) where {F, OP, A}
     if alg isa BlockReduce
         mapreduce_1d_gpu(
             f, op, src, backend;
             init, neutral,
             block_size=alg.block_size, items_per_thread=alg.items_per_thread,
+            max_blocks=reduce_tuning(backend, A).target_blocks,
             partials, switch_below=alg.switch_below,
         )
     else

src/reduce/tuning.jl

@@ -4,12 +4,15 @@
 """
     BlockReduce(; block_size=nothing, items_per_thread=nothing, switch_below=nothing)
 
-GPU tree reduction: each block of `block_size` threads (a power of two up to 1024) reduces
-`block_size * items_per_thread` elements in local memory, repeatedly, until one value is left.
-Reductions to a scalar (`dims=:`) finish on the host once fewer than `switch_below` values
-remain; for them, `block_size * items_per_thread` must be between 2 and `typemax(Int32)`. Reductions along `dims` choose among several launch shapes by the sizes and strides of
-the input and output; `items_per_thread` and `switch_below` do not apply to them, and setting
-either is an `ArgumentError` there.
+GPU tree reduction in blocks of `block_size` threads (a power of two up to 1024), which combine
+their values in local memory. A whole-array reduction (`dims=:`) takes at most two launches: in
+the first, each block reduces tiles of `block_size * items_per_thread` elements, one after the
+other, to a partial result, and a single block then reduces those partial results. Reductions to
+a scalar finish on the host instead once fewer than `switch_below` values remain; for them,
+`block_size * items_per_thread` must be between 2 and `typemax(Int32)`. Reductions along `dims`
+choose among several launch shapes by the sizes and strides of the input and output (into a
+single output, that of whole arrays); `items_per_thread` and `switch_below` cannot be set for
+them, and setting either is an `ArgumentError` there.
 
 Like every GPU reduction in the Julia ecosystem, it requires an associative and commutative
 operator: the combination order is fixed for a given setting and shape, but it is not the
@@ -28,20 +31,24 @@ end
 Values that fill unset algorithm fields for reductions on one device, as returned by
 [`reduce_tuning`](@ref).
 
-- `block_size`, `items_per_thread`, `switch_below`: settings for unset `BlockReduce` fields.
-- `target_blocks`: the number of blocks a reduction along `dims` aims to launch, to fill the
-  device when there are few outputs; must be positive.
+- `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 (see
+  `_reduce_items_per_thread`).
+- `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.
 - `columns_min_elements`: the fewest elements of a reduction along `dims` into contiguous outputs
   for which it splits each output's reduction across blocks of consecutive outputs; that can take
   a second launch, which costs more on some backends.
 - `threads_min_elems`: the default `min_elems` of `CPUThreads.Partitioned`.
 
-`columns_min_elements` defaults to 2^18, measured on CUDA; the other defaults reproduce AK's
-historical settings. Internal: the fields may change in any release.
+`items_per_thread` defaults to the choice by type, and `columns_min_elements` to 2^18, measured
+on CUDA; the other defaults reproduce AK's historical settings. Internal: the fields may change in
+any release.
 """
 Base.@kwdef struct ReduceTuning
     block_size::Int = 256
-    items_per_thread::Int = 2
+    items_per_thread::Union{Nothing, Int} = nothing
     switch_below::Int = 0
     target_blocks::Int = 256
     columns_min_elements::Int = 1 << 18
@@ -56,6 +63,15 @@ The reduction tuning for accumulator type `T` on `backend`'s current device; see
 """
 reduce_tuning(::Backend, ::Type) = ReduceTuning()
 
+# Elements each thread loads per tile of a whole-array reduction: enough for 64 bytes of
+# accumulators, from 4 to 16. The first pass launches at most `target_blocks` blocks, so its
+# memory-level parallelism comes from these loads, which each thread issues together. Measured
+# on an RTX 5080 (kernel time of `sum`, 256 blocks): 16 for 4-byte types is within 1% of the best
+# setting from 10^4 to 10^8 elements, while 2 reaches the bandwidth only with many more blocks,
+# and wider accumulators (Float64, ComplexF64, tuples of 4 Float64s) do best at 4 to 8.
+_reduce_items_per_thread(::Type{T}) where {T} =
+    isbitstype(T) && sizeof(T) > 0 ? clamp(64 ÷ sizeof(T), 4, 16) : 4
+
 
 _whole(dims) = dims === nothing || dims isa Colon
 
@@ -94,7 +110,7 @@ end
 
 _fill(a::BlockReduce, t::ReduceTuning, T) =
     BlockReduce(something(a.block_size, t.block_size),
-                something(a.items_per_thread, t.items_per_thread),
+                something(a.items_per_thread, t.items_per_thread, _reduce_items_per_thread(T)),
                 something(a.switch_below, t.switch_below))
 _fill(a::CPUThreads.Partitioned, t::ReduceTuning, T) = _fill_threads(a, t)
 _fill(a::Algorithm, t::ReduceTuning, T) =

src/reduce/utilities.jl

@@ -185,7 +185,7 @@ struct _Fold end
 @inline _finish(op, ::_NoInit, dst, i, partial) = _unlane(partial)
 @inline _finish(op, ::_Fold, dst, i, partial) = op(dst[i], _unlane(partial))
 
-# `init` of a pass that stores bare partial results, for a later pass
+# `init` of a pass that stores bare partial results, for a later pass or for the host
 struct _NoFinish end
 @inline _finish(op, ::_NoFinish, dst, i, partial) = partial
 

src/workspace.jl

@@ -81,7 +81,7 @@ julia> AK.workspace_size(AK.sort!, CuArray(rand(Float32, 10_000)); alg=AK.RadixS
  key_range = (partials = (Tuple{UInt32, UInt32}, (40,)),))
 
 julia> AK.workspace_size(AK.sum, CuArray(rand(Float32, 10_000)))
-(partials = (Float32, (40,)),)
+(partials = (Float32, (4,)),)
 ```
 """
 workspace_size(op, args...; kwargs...) = _public_sizes(_plan(op, args...; kwargs...).sizes)