diff --git a/src/reduce/mapreduce_1d_gpu.jl b/src/reduce/mapreduce_1d_gpu.jl index 73f1f0ac..48126b46 100644 --- a/src/reduce/mapreduce_1d_gpu.jl +++ b/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 diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index d35ec22f..bd49707b 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/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, diff --git a/src/reduce/reduce.jl b/src/reduce/reduce.jl index a79f976c..6f1bcf89 100644 --- a/src/reduce/reduce.jl +++ b/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 diff --git a/src/reduce/tuning.jl b/src/reduce/tuning.jl index 9f355b1a..7c9b9043 100644 --- a/src/reduce/tuning.jl +++ b/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) = diff --git a/src/reduce/utilities.jl b/src/reduce/utilities.jl index 8af4e074..6159d859 100644 --- a/src/reduce/utilities.jl +++ b/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 diff --git a/src/workspace.jl b/src/workspace.jl index 028f24e6..1d6fb246 100644 --- a/src/workspace.jl +++ b/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) diff --git a/test/generic/predicates.jl b/test/generic/predicates.jl index a6359b05..dc6ab150 100644 --- a/test/generic/predicates.jl +++ b/test/generic/predicates.jl @@ -82,10 +82,10 @@ struct PredicateCWTestBackend <: KernelAbstractions.Backend end AK.CPUThreads.Partitioned(Threads.nthreads(), 1) @test AK._resolve_predicate(AK.Auto(), PredicateCWTestBackend(), Float32) === AK.ConcurrentWrite(256) - @test AK._resolve_predicate(AK.Auto(), B, Float32) === AK.ViaReduce(AK.BlockReduce(256, 2, 0)) + @test AK._resolve_predicate(AK.Auto(), B, Float32) === AK.ViaReduce(AK.BlockReduce(256, 16, 0)) # The nested reduction is resolved like any reduction @test AK._resolve_predicate(AK.ViaReduce(AK.BlockReduce(block_size=64)), B, Float32) === - AK.ViaReduce(AK.BlockReduce(64, 2, 0)) + AK.ViaReduce(AK.BlockReduce(64, 16, 0)) @test_throws ArgumentError AK._resolve_predicate(AK.ViaReduce(AK.BlockReduce(block_size=3)), B, Float32) # A tuning only guides `Auto` @test AK._resolve_predicate(AK.ConcurrentWrite(), B, Float32) === AK.ConcurrentWrite(256) diff --git a/test/generic/reduce.jl b/test/generic/reduce.jl index c0d1e464..94873a40 100644 --- a/test/generic/reduce.jl +++ b/test/generic/reduce.jl @@ -1117,6 +1117,54 @@ end end +@testset "reductions into a single output" begin + # Along `dims`, a reduction into one output of a contiguous or fused source takes the + # whole-array kernels, finishing on the device; the sizes span one tile, one pass and two + Random.seed!(0) + alg = REDUCE_ALG + for n in (2, 1000, 50_000, 3_000_000) + h = rand(Int32(-9):Int32(9), n) + hm = reshape(h, :, 2) + v = array_from_host(h) + m = array_from_host(hm) + # `init` applied once, folding into `R`'s value, overwriting it + R = array_from_host(Int32[5]) + AK.mapreducedim!(identity, +, R, v; alg) + @test Array(R) == [5 + sum(h)] + AK.mapreducedim!(identity, +, R, v; init=Int32(1), alg) + @test Array(R) == [1 + sum(h)] + AK.mapreducedim!(identity, +, R, m; overwrite=true, alg) + @test Array(R) == [sum(h)] + # ... in the accumulator type, and without a neutral element + r = AK.sum(m; dims=(1, 2), alg) + @test eltype(r) === Int && Array(r) == sum(hm; dims=(1, 2)) + @test Array(AK.reduce((a, b) -> max(a, b), m; dims=(1, 2), alg)) == fill(maximum(h), 1, 1) + # A contiguous view at an offset, and a fused source + R = array_from_host(Int32[0]) + AK.mapreducedim!(identity, +, R, view(m, :, 2); alg) + @test Array(R) == [sum(view(hm, :, 2))] + @test Array(AK.mapreduce(*, +, m, m; dims=(1, 2), alg)) == fill(sum(h .* h), 1, 1) + # Partial results in the accumulator type, which only the result must fit `R`'s + n == 2 && continue + h8 = Int8[fill(Int8(100), n ÷ 2); fill(Int8(-100), n ÷ 2)] + h8[end] = -95 + v8 = array_from_host(h8) + R = array_from_host(Int8[0]) + AK.mapreducedim!(identity, +, R, v8; acctype=Int32, overwrite=true, alg) + @test Array(R) == Int8[5] + AK.mapreducedim!(identity, +, R, v8; acctype=Int32, init=Int8(-10), alg) + @test Array(R) == Int8[-5] + # ... with a workspace, used twice + R = array_from_host(Int8[0]) + ws = AK.workspace(AK.mapreducedim!, identity, +, R, v8; acctype=Int32, alg) + AK.mapreducedim!(identity, +, R, v8; acctype=Int32, alg, workspace=ws) + @test Array(R) == Int8[5] + AK.mapreducedim!(identity, +, R, v8; acctype=Int32, alg, workspace=ws) + @test Array(R) == Int8[10] + end +end + + @testset "reduction contract" begin Random.seed!(0) alg = REDUCE_ALG diff --git a/test/generic/reduce_resolve.jl b/test/generic/reduce_resolve.jl index 0d7f26ff..b3b76757 100644 --- a/test/generic/reduce_resolve.jl +++ b/test/generic/reduce_resolve.jl @@ -26,11 +26,11 @@ end @testset "reduce resolution: Auto" begin - # The default tuning reproduces AK's historical settings, for whole arrays and `dims` - @test resolve_reduce(AK.Auto()) === AK.BlockReduce(256, 2, 0) - @test resolve_reduce(AK.Auto(); dims=nothing) === AK.BlockReduce(256, 2, 0) - @test resolve_reduce(AK.Auto(); dims=2) === AK.BlockReduce(256, 2, 0) - @test resolve_reduce(AK.Auto(); dims=(1, 3)) === AK.BlockReduce(256, 2, 0) + # The default tuning, for whole arrays and `dims` + @test resolve_reduce(AK.Auto()) === AK.BlockReduce(256, 16, 0) + @test resolve_reduce(AK.Auto(); dims=nothing) === AK.BlockReduce(256, 16, 0) + @test resolve_reduce(AK.Auto(); dims=2) === AK.BlockReduce(256, 16, 0) + @test resolve_reduce(AK.Auto(); dims=(1, 3)) === AK.BlockReduce(256, 16, 0) # On the host backend: the threaded algorithm, filled host = AK.HOST_BACKEND @@ -44,6 +44,15 @@ end @test resolve_reduce(AK.Auto()) === AK.BlockReduce(512, 4, 100) @test resolve_reduce(AK.Auto(stable=false)) === AK.BlockReduce(512, 4, 100) end + + # A tuning without `items_per_thread` loads 64 bytes of accumulators per thread and tile, 4 to + # 16 elements + @test resolve_reduce(AK.Auto(); T=Float64) === AK.BlockReduce(256, 8, 0) + @test resolve_reduce(AK.Auto(); T=NTuple{8, Float64}) === AK.BlockReduce(256, 4, 0) + @test resolve_reduce(AK.Auto(); T=Union{}) === AK.BlockReduce(256, 4, 0) + with_reduce_tuning(; block_size=512, switch_below=100) do + @test resolve_reduce(AK.Auto(); T=Int8) === AK.BlockReduce(512, 16, 100) + end end @@ -83,7 +92,7 @@ end # and at most typemax(Int32), whether the settings are explicit or from the tuning @test_throws ArgumentError resolve_reduce(AK.BlockReduce(block_size=1, items_per_thread=1)) @test resolve_reduce(AK.BlockReduce(block_size=1, items_per_thread=2)) === AK.BlockReduce(1, 2, 0) - @test resolve_reduce(AK.BlockReduce(block_size=1); dims=1) === AK.BlockReduce(1, 2, 0) + @test resolve_reduce(AK.BlockReduce(block_size=1); dims=1) === AK.BlockReduce(1, 16, 0) @test_throws ArgumentError resolve_reduce(AK.BlockReduce(items_per_thread=1 << 56)) @test_throws ArgumentError resolve_reduce(AK.BlockReduce(block_size=1024, items_per_thread=1 << 22)) with_reduce_tuning(; block_size=1, items_per_thread=1) do @@ -95,7 +104,7 @@ end @test_throws ArgumentError resolve_reduce(AK.CPUThreads.Partitioned()) @test_throws ArgumentError resolve_reduce(AK.BlockReduce(); backend=ReduceNoKernelsTestBackend()) if AK._runs_kernels(host) - @test resolve_reduce(AK.BlockReduce(); backend=host) === AK.BlockReduce(256, 2, 0) + @test resolve_reduce(AK.BlockReduce(); backend=host) === AK.BlockReduce(256, 16, 0) else @test_throws ArgumentError resolve_reduce(AK.BlockReduce(); backend=host) end diff --git a/test/generic/workspace.jl b/test/generic/workspace.jl index 52a18b51..f7aa2fb0 100644 --- a/test/generic/workspace.jl +++ b/test/generic/workspace.jl @@ -51,8 +51,18 @@ @test Base.all(k -> haskey(sizes, k), (:temp, :hist, :scan, :key_range)) @test sizes.temp == (Float32, (10_000,)) @test AK.workspace_size(AK.sort!, v; alg=AK.BitonicSort()) == (;) - # A whole-array reduction keeps its partial results in two halves of one buffer + # A whole-array reduction keeps the partial results of its first pass, at most + # `target_blocks`, and a slot for the last one @test haskey(AK.workspace_size(AK.sum, v; alg=AK.BlockReduce()), :partials) + @test AK.workspace_size(AK.sum, array_from_host(rand(Float32, 10^7)); + alg=AK.BlockReduce()).partials[2] == + (AK.reduce_tuning(BACKEND, Float32).target_blocks + 1,) + # ... as does one along `dims` into a single output, which has no slot + @test AK.workspace_size(AK.mapreducedim!, identity, +, + array_from_host(zeros(Float32, 1, 1)), + array_from_host(rand(Float32, 1000, 10_000)); + alg=AK.BlockReduce()).partials[2] == + (AK.reduce_tuning(BACKEND, Float32).target_blocks,) # ... which a reduction of one element does not need @test AK.workspace_size(AK.sum, v[1:1]; alg=AK.BlockReduce()) == (; partials=(Float32, (0,))) # DecoupledLookback's flags only where it runs @@ -155,11 +165,12 @@ end # ... and a reduction against every array of a fused source x = array_from_host(ones(Int32, 1000)) y = copy(x) + # (two elements in one tile: the partials are the first pass's one and the result's) a = AK.BlockReduce(block_size=2, items_per_thread=1) - ws = AK.workspace(AK.mapreduce, +, +, x, y; alg=a) + ws = AK.workspace(AK.mapreduce, +, +, x[1:2], y[1:2]; alg=a) z = ws.buffers.partials - @test length(z) == 1000 - @test_throws ArgumentError AK.mapreduce(+, +, z, y; alg=a, workspace=ws) + @test length(z) == 2 + @test_throws ArgumentError AK.mapreduce(+, +, z, y[1:2]; alg=a, workspace=ws) # ... and a findall against its input ws = AK.workspace(AK.findall, identity, array_from_host(rand(Bool, 1000))) @test_throws ArgumentError AK.findall(identity, ws.buffers.mask; workspace=ws)