diff --git a/ext/oneAPIExt.jl b/ext/oneAPIExt.jl index fa5bdfd0..059a9873 100644 --- a/ext/oneAPIExt.jl +++ b/ext/oneAPIExt.jl @@ -17,5 +17,10 @@ oneAPI.@device_override AK._decoupled_fence() = # keeps using `ViaReduce` for `any`/`all` on oneAPI; an explicit `ConcurrentWrite` is allowed. AK.predicate_tuning(::oneAPIBackend, ::Type) = AK.PredicateTuning(prefer_concurrent_write=false) +# Launches cost tens of microseconds more than on other GPUs (Iris Xe: `sum(A; dims=2)` of a +# 64×4096 matrix takes 220 µs in one launch and 280 µs in the column reduction's two; from about +# 4M elements the column reduction is faster), so the column reduction starts at larger inputs. +AK.reduce_tuning(::oneAPIBackend, ::Type) = AK.ReduceTuning(columns_min_elements=1 << 22) + end # module oneAPIExt diff --git a/src/reduce/mapreduce_1d_cpu.jl b/src/reduce/mapreduce_1d_cpu.jl index 3d31ed10..d6955563 100644 --- a/src/reduce/mapreduce_1d_cpu.jl +++ b/src/reduce/mapreduce_1d_cpu.jl @@ -1,12 +1,12 @@ # Reduce the non-empty `src` to a value on Julia threads; `neutral` and `init` as for # `mapreduce_1d_gpu`. function mapreduce_1d_cpu( - f, op, src::MapReduceSource, backend::Backend; + f::F, op::OP, src::MapReduceSource, backend::Backend; init, neutral, max_tasks::Int, min_elems::Int, -) +) where {F, OP} f, op_lanes = _lanefuncs(f, op, neutral) tp = TaskPartitioner(length(src), max_tasks, min_elems) if src isa Base.Broadcast.Broadcasted || tp.num_tasks == 1 diff --git a/src/reduce/mapreduce_1d_gpu.jl b/src/reduce/mapreduce_1d_gpu.jl index 51d27883..73f1f0ac 100644 --- a/src/reduce/mapreduce_1d_gpu.jl +++ b/src/reduce/mapreduce_1d_gpu.jl @@ -42,14 +42,14 @@ end # `_reduce_seed`, `init` is a value or `_NoInit()`, and `partials` holds `_mapreduce_1d_partials` # partial results (`nothing` if none are needed). function mapreduce_1d_gpu( - f, op, src::MapReduceSource, backend::Backend; + f::F, op::OP, src::MapReduceSource, backend::Backend; init, neutral, block_size::Int, items_per_thread::Int, partials::Union{Nothing, AbstractArray}, switch_below::Int, -) +) where {F, OP} @argcheck 1 <= block_size <= 1024 @argcheck ispow2(block_size) @argcheck items_per_thread >= 1 diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index 0960ecc8..d35ec22f 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -8,9 +8,11 @@ # 2. The inner reduce loop must avoid per-element integer division. For a single segment the # element offset is just `j * stride`; only genuinely non-contiguous dim sets (e.g. # `dims=(1,3)`) fall back to a per-element multi-dimensional decode. -# 3. Four work decompositions, chosen by the relative sizes and strides of the output and +# 3. Five work decompositions, chosen by the relative sizes and strides of the output and # the reduction: # - by_thread: one thread per output (many outputs, small reduction) +# - columns: 32 contiguous outputs per block, split across blocks, for long strided +# reductions (e.g. `dims=2` of a wide matrix) # - tiled_strided: several contiguous outputs per block, for square-like strided reductions # - by_block: one block per output, grid-stride over outputs (few outputs, large reduction) # - multigroup: several blocks per output, two-pass (dst_size==1 or very small dst_size) @@ -29,54 +31,81 @@ const MIN_ITEMS_PER_THREAD = 8 # 256-thread block this gives 32 lanes per row. const TILED_STRIDED_ROWS_PER_BLOCK = 8 +# The most consecutive outputs per block of the columns kernel (fewer outputs take the next power +# of two), so that a warp's loads of one step are consecutive elements; and the loads each of its +# threads issues together +const COLUMNS_COLS = 32 +const COLUMNS_ITEMS = 8 +# The fewest outputs the columns kernel is used for: with 4 or fewer, its blocks' scattered loads +# lose to `:multigroup` on Metal and oneAPI +const COLUMNS_MIN_OUTPUTS = 8 +# The most threads per block of the columns kernel: its unrolled loads take more registers than +# the other kernels, and Metal allows it fewer than 1024 threads per block +const COLUMNS_MAX_THREADS = 256 -# Host-side canonicalization: split the dimensions into reduced and kept ("outer") segments, -# merging adjacent dimensions whose strides are contiguous. Returns two tuples of -# (stride, size) segments. -function _canonicalize_dims(src_sizes, src_strides, dims_valid) - ndim = length(src_sizes) - reduce_segs = Tuple{Int,Int}[] - outer_segs = Tuple{Int,Int}[] +# Host-side canonicalization: split the dimensions into reduced and kept ("outer") segments, +# merging adjacent dimensions whose strides are contiguous. Returns the `(stride, size)` pairs of +# the reduced and of the kept segments, each as `N`-tuples of strides and sizes padded with +# zeros, and the number of segments. (The fixed length keeps this type stable: the reduction's +# host path runs on every call, and dynamic dispatch there costs more than small kernels take.) +function _canonicalize_dims(src_sizes::NTuple{N, Int}, src_strides::NTuple{N, Int}, + dims_valid) where {N} + pad = ntuple(_ -> 0, Val(N)) + rstrides, rsizes, nr = pad, pad, 0 + ostrides, osizes, no = pad, pad, 0 i = 1 - while i <= ndim - if i in dims_valid - seg_stride = src_strides[i] - seg_size = src_sizes[i] - while i + 1 <= ndim && - (i + 1) in dims_valid && - src_strides[i + 1] == seg_stride * seg_size - i += 1 - seg_size *= src_sizes[i] - end - push!(reduce_segs, (seg_stride, seg_size)) + while i <= N + reduced = i in dims_valid + seg_stride = src_strides[i] + seg_size = src_sizes[i] + while i + 1 <= N && ((i + 1) in dims_valid) == reduced && + src_strides[i + 1] == seg_stride * seg_size + i += 1 + seg_size *= src_sizes[i] + end + if reduced + nr += 1 + rstrides = Base.setindex(rstrides, seg_stride, nr) + rsizes = Base.setindex(rsizes, seg_size, nr) else - seg_stride = src_strides[i] - seg_size = src_sizes[i] - while i + 1 <= ndim && - !((i + 1) in dims_valid) && - src_strides[i + 1] == seg_stride * seg_size - i += 1 - seg_size *= src_sizes[i] - end - push!(outer_segs, (seg_stride, seg_size)) + no += 1 + ostrides = Base.setindex(ostrides, seg_stride, no) + osizes = Base.setindex(osizes, seg_size, no) end i += 1 end - - return Tuple(reduce_segs), Tuple(outer_segs) + return (rstrides, rsizes, nr), (ostrides, osizes, no) end -# The reduced and kept dimensions of a strided source, as `(strides, sizes)` of their merged -# segments. -function _mapreduce_segments(src_sizes, src_strides, dims_valid) - reduce_segs, outer_segs = _canonicalize_dims(src_sizes, src_strides, dims_valid) - return (Tuple(str for (str, _) in reduce_segs), Tuple(s for (_, s) in reduce_segs)), - (Tuple(str for (str, _) in outer_segs), Tuple(s for (_, s) in outer_segs)) +# The reduced and kept dimensions of a strided source, as `(strides, sizes, count)` of their +# merged segments (see `_canonicalize_dims`) +_mapreduce_segments(src_sizes, src_strides, dims_valid) = + _canonicalize_dims(Tuple(src_sizes), Tuple(src_strides), dims_valid) + +# Call `g(reduce_strides, reduce_sizes, outer_strides, outer_sizes)` with the segments of +# `_mapreduce_segments` as tuples of their lengths, which the kernels specialize on. The common +# layouts (one reduced segment and at most one kept one) take no dynamic dispatch. +@inline function _with_segments(g::G, segs) where {G} + (rstrides, rsizes, nr), (ostrides, osizes, no) = segs + if nr == 1 && no == 1 + return g((rstrides[1],), (rsizes[1],), (ostrides[1],), (osizes[1],)) + elseif nr == 1 && no == 0 + return g((rstrides[1],), (rsizes[1],), (), ()) + else + return g(_segment_tuple(rstrides, nr), _segment_tuple(rsizes, nr), + _segment_tuple(ostrides, no), _segment_tuple(osizes, no)) + end end +_segment_tuple(t::Tuple, n::Int) = ntuple(i -> t[i], n) # The launch shape of a strided reduction along `dims`, and the number of blocks per output of a -# `:multigroup` reduction: +# `:multigroup` or `:columns` reduction: +# - at least COLUMNS_MIN_OUTPUTS contiguous outputs +# with a longer strided reduction, of at least +# the tuning's `columns_min_elements` in all -> :columns (blocks of up to 32 +# consecutive outputs; a second pass +# when an output's reduction is split) # - square-like, contiguous output with strided input -> :tiled_strided # - dst_size >= reduce_size -> :by_thread # - dst_size == 1, or dst_size < GS_DST_CUTOFF @@ -92,17 +121,17 @@ end # the (large) reduction itself across many blocks via multigroup is still better. # For dst_size==1 there is nothing to grid-stride over, so multigroup is the only # option regardless of GS_DST_CUTOFF. -function _mapreduce_nd_shape( - dst_size, reduce_size, reduce_strides, reduce_sizes, outer_strides, outer_sizes, - block_size, target_blocks, -) +function _mapreduce_nd_shape(dst_size, reduce_size, segs, block_size, t::ReduceTuning) + target_blocks = t.target_blocks + (rstrides, rsizes, nr), (ostrides, _, no) = segs + single_reduce = nr == 1 && rsizes[1] != 0 + # by_block override 1: when the reduced dimension is the fastest-varying one # (reduce_strides==(1,)), by_thread's per-thread strided access is badly # uncoalesced, while by_block lets consecutive threads read consecutive elements. use_by_block_for_coalescing = dst_size >= reduce_size && reduce_size >= block_size && - length(reduce_sizes) == 1 && reduce_sizes[1] != 0 && - reduce_strides == (1,) + single_reduce && rstrides[1] == 1 # by_block override 2: when by_thread would launch too few blocks for a square-like # output/reduction shape, split each output reduction across a full block. This @@ -112,17 +141,28 @@ function _mapreduce_nd_shape( # coalescing is better than a strided block reduction. use_by_block_for_low_occupancy = dst_size == reduce_size && reduce_size >= block_size && - cld(dst_size, block_size) < target_blocks && - length(reduce_sizes) == 1 && reduce_sizes[1] != 0 + cld(dst_size, block_size) < target_blocks && single_reduce use_by_block = use_by_block_for_coalescing || use_by_block_for_low_occupancy use_tiled_strided = dst_size == reduce_size && reduce_size >= block_size && - length(outer_sizes) == 1 && outer_strides == (1,) && - length(reduce_sizes) == 1 && reduce_strides[1] > 1 && + no == 1 && ostrides[1] == 1 && nr == 1 && rstrides[1] > 1 && block_size % TILED_STRIDED_ROWS_PER_BLOCK == 0 && ispow2(block_size ÷ TILED_STRIDED_ROWS_PER_BLOCK) + # Contiguous outputs with a longer strided reduction, e.g. the row sums of a wide matrix: + # :columns, with enough blocks per output to fill the device + if no == 1 && ostrides[1] == 1 && nr == 1 && rstrides[1] > 1 && + COLUMNS_MIN_OUTPUTS <= dst_size < reduce_size && + widemul(dst_size, reduce_size) >= t.columns_min_elements && + block_size >= _columns_cols(dst_size) + cols = _columns_cols(dst_size) + lanes = min(block_size, COLUMNS_MAX_THREADS) ÷ cols + groups = clamp(cld(target_blocks, cld(dst_size, cols)), 1, + cld(reduce_size, lanes * MIN_ITEMS_PER_THREAD)) + return :columns, groups + end + if use_tiled_strided return :tiled_strided, 0 elseif dst_size >= reduce_size && !use_by_block @@ -142,8 +182,12 @@ function _mapreduce_nd_shape( end end +# The outputs per block of the columns kernel +_columns_cols(dst_size) = min(nextpow(2, dst_size), COLUMNS_COLS) + # The scratch of a reduction along `dims` into `dst_size` outputs: the partial results of a -# `:multigroup` reduction, with the partial-result type of the seed `neutral` (see `_reduce_seed`) +# `:multigroup` reduction or of a `:columns` one in two passes, 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) @@ -151,14 +195,10 @@ function _mapreduce_nd_sizes(src, backend, alg, ::Type{A}, neutral, dims_valid, (len <= 1 || dst_size == 0) && return (;) layout = _mapreduce_strided_layout(src) isnothing(layout) && return (;) - target_blocks = reduce_tuning(backend, A).target_blocks - (reduce_strides, reduce_sizes), (outer_strides, outer_sizes) = - _mapreduce_segments(src_sizes, layout[3], dims_valid) - shape, reduce_groups = _mapreduce_nd_shape( - dst_size, len, reduce_strides, reduce_sizes, outer_strides, outer_sizes, - alg.block_size, target_blocks, - ) - shape === :multigroup || return (;) + segs = _mapreduce_segments(src_sizes, layout[3], dims_valid) + shape, reduce_groups = _mapreduce_nd_shape(dst_size, len, segs, alg.block_size, + reduce_tuning(backend, A)) + shape === :multigroup || shape === :columns && reduce_groups > 1 || return (;) return (; partials=_buffer(typeof(neutral), dst_size, reduce_groups)) end @@ -182,13 +222,13 @@ end # dimensions `dims_valid`; `A` is the accumulator type, and `init` is a value, `_NoInit()` or # `_Fold()` (see `_finish`). function mapreduce_nd!( - f, op, dst::AbstractArray, src::MapReduceSource, backend::Backend, + f::F, op::OP, dst::AbstractArray, src::MapReduceSource, backend::Backend, alg::Union{BlockReduce, CPUThreads.Partitioned}, ::Type{A}; init, neutral, dims_valid::Tuple, bufs::NamedTuple, -) where {A} +) where {F, OP, A} # Launch settings of the elementwise passes, and of the reduction kernels launch = alg isa BlockReduce ? (; block_size=alg.block_size) : (; max_tasks=alg.max_tasks, min_elems=alg.min_elems) @@ -196,7 +236,8 @@ function mapreduce_nd!( # Number of blocks the by_block / multigroup paths aim to launch, so a reduction with # few output elements can still fill the GPU. A heuristic GPU-occupancy target. - target_blocks = reduce_tuning(backend, A).target_blocks + tuning = reduce_tuning(backend, A) + target_blocks = tuning.target_blocks target_blocks >= 1 || throw(ArgumentError( "the reduction tuning's `target_blocks` must be positive, got $target_blocks")) @@ -253,21 +294,52 @@ function mapreduce_nd!( end buffer, base_offset, src_strides = layout - (reduce_strides, reduce_sizes), (outer_strides, outer_sizes) = - _mapreduce_segments(src_sizes, src_strides, dims_valid) + segs = _mapreduce_segments(src_sizes, src_strides, dims_valid) reduce_size = len - shape, reduce_groups = _mapreduce_nd_shape( - dst_size, reduce_size, reduce_strides, reduce_sizes, outer_strides, outer_sizes, - block_size, target_blocks, - ) + shape, reduce_groups = _mapreduce_nd_shape(dst_size, reduce_size, segs, block_size, tuning) + _with_segments(segs) do reduce_strides, reduce_sizes, outer_strides, outer_sizes + _mapreduce_nd_launch!( + shape, reduce_groups, f, op, dst, buffer, backend, block_size, target_blocks; + init, neutral, bufs, base_offset, reduce_strides, reduce_sizes, outer_strides, + outer_sizes, dst_size, reduce_size, + ) + end + + return dst +end - if shape === :tiled_strided +# Launch the kernels of a strided reduction of shape `shape` (see `_mapreduce_nd_shape`) +function _mapreduce_nd_launch!( + shape, reduce_groups, f::F, op::OP, dst, buffer, backend, block_size, target_blocks; + init, neutral, bufs, base_offset, reduce_strides, reduce_sizes, outer_strides, outer_sizes, + dst_size, reduce_size, +) where {F, OP} + if shape === :columns + groups = reduce_groups + blocks = cld(dst_size, _columns_cols(dst_size)) * groups + threads = min(block_size, COLUMNS_MAX_THREADS) + kernel! = _mapreduce_nd_columns!(backend, threads) + cols = Val(_columns_cols(dst_size)) + if groups == 1 + kernel!(buffer, dst, f, op, neutral, init, base_offset, reduce_strides[1], dst_size, + reduce_size, 1, cols; ndrange=(threads * blocks,)) + else + # (the partial results of each output's groups, combined by a second pass) + partial = bufs.partials + kernel!(buffer, partial, f, op, neutral, _NoFinish(), base_offset, + reduce_strides[1], dst_size, reduce_size, groups, cols; + ndrange=(threads * blocks,)) + pass2_block_size = _pass2_block_size(groups) + kernel2! = _mapreduce_partial_to_dst!(backend, pass2_block_size) + kernel2!(partial, dst, op, neutral, init, dst_size, groups, + ndrange=(pass2_block_size * dst_size,)) + end + elseif shape === :tiled_strided # Narrow layout-specific path: square-like row reductions with contiguous # outputs and strided input reads, e.g. size=(1024,1024), dims=2. Several # outputs share a block: small lane groups reduce one output each, preserving # more cross-output memory coalescing than by_block while launching many more - # blocks than by_thread. This is the only extra kernel beyond the generic - # by_thread/by_block/multigroup shapes. + # blocks than by_thread. rows_per_block = TILED_STRIDED_ROWS_PER_BLOCK blocks = cld(dst_size, rows_per_block) kernel! = _mapreduce_nd_by_thread_tiled_strided!(backend, block_size) @@ -320,17 +392,14 @@ function mapreduce_nd!( ) end - return dst + return nothing end -function _mapreduce_dense_strides(sizes::Tuple) - stride = 1 - return ntuple(length(sizes)) do i - s = stride - stride *= sizes[i] - s - end -end +# (the strides of a dense column-major array: 1 and the running products of the sizes) +_mapreduce_dense_strides(sizes::Tuple{}) = () +_mapreduce_dense_strides(sizes::Tuple) = Base.front(_cumprod((1, sizes...))) +_cumprod(t::Tuple{Any}) = t +_cumprod(t::Tuple) = (t[1], _cumprod((t[1] * t[2], Base.tail(Base.tail(t))...))...) # Resolve the layout the stride-based fast-path kernels need. Those kernels index a # flat buffer at `buffer[base_offset + Σ coordᵈ·strideᵈ + 1]`, so they work for any @@ -420,10 +489,10 @@ function _contiguous_vector_width( end function _launch_mapreduce_nd_by_block!( - backend, block_size, buffer, dst, f, op, neutral, init, + backend, block_size, buffer, dst, f::F, op::OP, neutral, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, output_size, reduce_size, num_blocks, -) +) where {F, OP} W = _contiguous_vector_width( buffer, base_offset, outer_strides, reduce_strides, reduce_size, block_size, ) @@ -475,9 +544,9 @@ end # CPU path function _mapreduce_nd_cpu_sections!( - f, op, dst, src; + f::F, op::OP, dst, src; init, neutral, max_tasks, min_elems, -) +) where {F, OP} Rother = CartesianIndices(dst) Rreduce = _mapreduce_reduce_indices(src, dst) f_lanes, op_lanes = _lanefuncs(f, op, neutral) @@ -619,6 +688,73 @@ end end end +# GPU kernel: columns — a long strided reduction into contiguous outputs, e.g. the row sums of +# a wide column-major matrix (`dims=2`). Each block covers `cols` consecutive outputs with +# `N ÷ cols` lanes each, so that a warp's loads of one step are consecutive elements, and the +# reduction of each output is split across `groups` blocks. With one group the block stores the +# outputs; with more, `dst` holds `output_size × groups` partial results for +# `_mapreduce_partial_to_dst!` (and `init` is `_NoFinish()`). +@kernel inbounds=true cpu=false unsafe_indices=true function _mapreduce_nd_columns!( + @Const(src), dst, + f, op, neutral, init, + base_offset, reduce_stride, + output_size, reduce_size, groups, + ::Val{cols}, +) where {cols} + @uniform N = @groupsize()[1] + @uniform lanes = N ÷ cols + sdata = @localmem typeof(neutral) (N,) + f_lanes, op_lanes = _lanefuncs(f, op, neutral) + + iblock = @index(Group, Linear) - 0x1 + ithread = @index(Local, Linear) - 0x1 + tiles = cld(output_size, cols) + + # (unsigned divisions by constants and by the tile count; see the tiled kernel) + col = Int(unsigned(ithread) % unsigned(cols)) + lane = Int(unsigned(ithread) ÷ unsigned(cols)) + tile = Int(unsigned(iblock) % unsigned(tiles)) + group = Int(unsigned(iblock) ÷ unsigned(tiles)) + iout = tile * cols + col + + acc = neutral + if iout < output_size + sbase = base_offset + iout + step = groups * lanes + j = group * lanes + lane + # (several loads in flight per thread: one at a time leaves the kernel at half the + # bandwidth) + while j + (COLUMNS_ITEMS - 1) * step < reduce_size + for k in 0:(COLUMNS_ITEMS - 1) + acc = op_lanes(acc, f_lanes(src[sbase + (j + k * step) * reduce_stride + 0x1])) + end + j += COLUMNS_ITEMS * step + end + while j < reduce_size + acc = op_lanes(acc, f_lanes(src[sbase + j * reduce_stride + 0x1])) + j += step + end + end + + sdata[ithread + 0x1] = acc + @synchronize() + + s = lanes ÷ 2 + while s >= 1 + if lane < s + sdata[col + lane * cols + 0x1] = + op_lanes(sdata[col + lane * cols + 0x1], sdata[col + (lane + s) * cols + 0x1]) + end + @synchronize() + s ÷= 2 + end + + if lane == 0x0 && iout < output_size + i = iout + group * output_size + 0x1 + dst[i] = _finish(op, init, dst, i, sdata[col + 0x1]) + end +end + # GPU kernel: by_thread — one thread per output element, reducing sequentially. # Used when there are more output elements than elements in the reduced dimension(s), # e.g. reduce(+, rand(3, 1000), dims=1) — only 3 elements to reduce per output. diff --git a/src/reduce/reduce.jl b/src/reduce/reduce.jl index 9ab65927..a79f976c 100644 --- a/src/reduce/reduce.jl +++ b/src/reduce/reduce.jl @@ -37,7 +37,7 @@ AK.reduce(max, m; dims=1) # a 1×100_000 MtlArray AK.reduce(+, m; dims=2, alg=AK.BlockReduce(block_size=512)) ``` """ -reduce(op, src::AbstractArray; kwargs...) = mapreduce(identity, op, src; kwargs...) +reduce(op::OP, src::AbstractArray; kwargs...) where {OP} = mapreduce(identity, op, src; kwargs...) """ @@ -116,7 +116,7 @@ rows = AK.mapreduce((x, y) -> x * y, +, a, b; dims=1) ``` """ function mapreduce( - f, op, src::MapReduceSource, srcs::AbstractArray...; + f::F, op::OP, src::MapReduceSource, srcs::AbstractArray...; backend::Union{Nothing, Backend}=nothing, init=_NoInit(), neutral=nothing, @@ -124,7 +124,7 @@ function mapreduce( dims=:, alg::Algorithm=Auto(), workspace=nothing, -) +) where {F, OP} f, src = _mapreduce_fuse(f, src, srcs) s = _mapreduce_setup(f, op, src, backend, init, neutral, acctype, dims, alg) bufs = _buffers(s.plan, workspace, src) @@ -132,20 +132,20 @@ function mapreduce( end function _plan( - ::typeof(mapreduce), f, op, src::MapReduceSource, srcs::AbstractArray...; + ::typeof(mapreduce), f::F, op::OP, src::MapReduceSource, srcs::AbstractArray...; backend=nothing, init=_NoInit(), neutral=nothing, acctype=nothing, dims=:, alg::Algorithm=Auto(), -) +) where {F, OP} f, src = _mapreduce_fuse(f, src, srcs) return _mapreduce_setup(f, op, src, backend, init, neutral, acctype, dims, alg).plan end -_plan(::typeof(reduce), op, src::AbstractArray; kwargs...) = +_plan(::typeof(reduce), op::OP, src::AbstractArray; kwargs...) where {OP} = _plan(mapreduce, identity, op, src; kwargs...) # The source of a reduction: several arrays become one `Broadcasted` object -_mapreduce_fuse(f, src, ::Tuple{}) = (f, src) -function _mapreduce_fuse(f, src, srcs::Tuple) +_mapreduce_fuse(f::F, src, ::Tuple{}) where {F} = (f, src) +function _mapreduce_fuse(f::F, src, srcs::Tuple) where {F} src isa AbstractArray || throw(ArgumentError("a Broadcasted source cannot be combined with more arrays")) _mapreduce_check_map_axes(src, srcs...) @@ -160,12 +160,13 @@ _mapreduce_source(src::Base.Broadcast.Broadcasted) = VERSION < v"1.12-" ? _materialize_source(src) : src # The seed of the partial results (`_reduce_seed`), where an accumulator type exists -_mapreduce_seed(op, ::Type{A}, neutral) where {A} = +_mapreduce_seed(op::OP, ::Type{A}, neutral) where {OP, A} = A === Union{} ? nothing : _reduce_seed(op, A, neutral) # Everything a reduction resolves before touching data: its backend, algorithm and scratch (the # plan), accumulator type and partial-result seed, and for reductions along `dims` the result -function _mapreduce_setup(f, op, src, backend, init, neutral, acctype, dims, alg) +function _mapreduce_setup(f::F, op::OP, src, backend, init, neutral, acctype, dims, + alg) where {F, OP} backend = _resolve_backend(backend, src) M = _mapped_eltype(f, src) # The accumulator type, from `init`'s type @@ -198,7 +199,7 @@ function _mapreduce_setup(f, op, src, backend, init, neutral, acctype, dims, alg end # Run the reduction set up by `_mapreduce_setup`, with the plan's scratch buffers `bufs` -function _mapreduce_run(f, op, src, s, init, bufs) +function _mapreduce_run(f::F, op::OP, src, s, init, bufs) where {F, OP} backend, a = s.plan.backend, s.plan.alg if !haskey(s, :dims_valid) return _mapreduce_whole(f, op, src, backend, a, _unval(s.A); init, neutral=s.seed, @@ -215,7 +216,8 @@ function _mapreduce_run(f, op, src, s, init, bufs) end # A reduction nested in another operation, with the scratch buffers of its plan in the outer one -_mapreduce_nested(f, op, src, bufs; backend, init=_NoInit(), neutral=nothing, dims=:, alg) = +_mapreduce_nested(f::F, op::OP, src, bufs; backend, init=_NoInit(), neutral=nothing, dims=:, + alg) where {F, OP} = _mapreduce_run(f, op, src, _mapreduce_setup(f, op, src, backend, init, neutral, nothing, dims, alg), init, bufs) @@ -284,7 +286,7 @@ AK.mapreducedim!(identity, +, R, A; init=100) # [103 107] ``` """ function mapreducedim!( - f, op, R::AbstractArray, src::MapReduceSource; + f::F, op::OP, R::AbstractArray, src::MapReduceSource; backend::Union{Nothing, Backend}=nothing, init=_NoInit(), neutral=nothing, @@ -292,20 +294,26 @@ function mapreducedim!( acctype=nothing, alg::Algorithm=Auto(), workspace=nothing, -) +) where {F, OP} s = _mapreducedim_setup(f, op, R, src, backend, init, neutral, overwrite, acctype, alg) + _mapreducedim_run!(f, op, R, src, s, workspace) + return R +end + +# (a function barrier: the type of `s` depends on the number of reduced dimensions) +function _mapreducedim_run!(f::F, op::OP, R, src, s, workspace) where {F, OP} bufs = _buffers(s.plan, workspace, R, src) mapreduce_nd!(f, op, s.dst, _mapreduce_source(s.src), s.plan.backend, s.plan.alg, _unval(s.A); init=s.init, neutral=s.seed, dims_valid=s.dims_valid, bufs) - return R end -_plan(::typeof(mapreducedim!), f, op, R::AbstractArray, src::MapReduceSource; +_plan(::typeof(mapreducedim!), f::F, op::OP, R::AbstractArray, src::MapReduceSource; backend=nothing, init=_NoInit(), neutral=nothing, overwrite::Bool=false, acctype=nothing, - alg::Algorithm=Auto()) = + alg::Algorithm=Auto()) where {F, OP} = _mapreducedim_setup(f, op, R, src, backend, init, neutral, overwrite, acctype, alg).plan -function _mapreducedim_setup(f, op, R, src, backend, init, neutral, overwrite, acctype, alg) +function _mapreducedim_setup(f::F, op::OP, R, src, backend, init, neutral, overwrite, acctype, + alg) where {F, OP} backend = _resolve_backend(backend, R, src) nd = ndims(src) for d in 1:max(nd, ndims(R)) @@ -317,6 +325,11 @@ function _mapreducedim_setup(f, op, R, src, backend, init, neutral, overwrite, a dst_sizes = ntuple(d -> size(R, d), nd) dst = size(R) == dst_sizes ? R : reshape(R, dst_sizes) dims_valid = Tuple(d for d in 1:nd if dst_sizes[d] == 1 && _srcsize(src, d) != 1) + return _mapreducedim_setup(f, op, src, dst, dims_valid, backend, init, neutral, overwrite, + acctype, alg) +end +function _mapreducedim_setup(f::F, op::OP, src, dst, dims_valid, backend, init, neutral, + overwrite, acctype, alg) where {F, OP} A = _acctype(op, eltype(dst), _mapped_eltype(f, src), acctype) a = _resolve_reduce(alg, backend, A, dims_valid) seed = _mapreduce_seed(op, A, neutral) @@ -336,7 +349,8 @@ _check_noalias(R, src) = nothing # Reduce all of `src` to a host value, with accumulator type `A` and partial-result seed # `neutral`; `partials` is the plan's scratch -function _mapreduce_whole(f, op, src, backend, alg, ::Type{A}; init, neutral, partials) where {A} +function _mapreduce_whole(f::F, op::OP, src, backend, alg, ::Type{A}; init, neutral, + partials) where {F, OP, A} if length(src) == 0 init isa _NoInit || return init throw(ArgumentError( @@ -352,7 +366,8 @@ function _mapreduce_whole(f, op, src, backend, alg, ::Type{A}; init, neutral, pa return convert(A, _mapreduce_whole_run(f, op, src, backend, alg; init, neutral, partials)) end -function _mapreduce_whole_run(f, op, src, backend, alg; init, neutral, partials) +function _mapreduce_whole_run(f::F, op::OP, src, backend, alg; init, neutral, + partials) where {F, OP} if alg isa BlockReduce mapreduce_1d_gpu( f, op, src, backend; diff --git a/src/reduce/tuning.jl b/src/reduce/tuning.jl index 25f081f8..9f355b1a 100644 --- a/src/reduce/tuning.jl +++ b/src/reduce/tuning.jl @@ -31,15 +31,20 @@ 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. - `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. +- `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`. -The defaults reproduce AK's historical settings. Internal: the fields may change in any release. +`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. """ Base.@kwdef struct ReduceTuning block_size::Int = 256 items_per_thread::Int = 2 switch_below::Int = 0 target_blocks::Int = 256 + columns_min_elements::Int = 1 << 18 threads_min_elems::Int = 1 end diff --git a/src/reduce/utilities.jl b/src/reduce/utilities.jl index d7cf0b6a..8af4e074 100644 --- a/src/reduce/utilities.jl +++ b/src/reduce/utilities.jl @@ -9,7 +9,7 @@ end # The mapped element type of a reduction source. Base's `combine_eltypes` does not see through # the `Extruded` wrappers of a preprocessed `Broadcasted`, so its element types are derived here. -_mapped_eltype(f, src) = Base.promote_op(f, _source_eltype(src)) +_mapped_eltype(f::F, src) where {F} = Base.promote_op(f, _source_eltype(src)) _source_eltype(src::AbstractArray) = eltype(src) _source_eltype(bc::Base.Broadcast.Broadcasted) = Base.promote_op(bc.f, Base.map(_source_eltype, bc.args)...) _source_eltype(x::Base.Broadcast.Extruded) = eltype(x.x) @@ -19,7 +19,8 @@ _source_eltype(x) = Base.Broadcast._broadcast_getindex_eltype(x) _materialize_source(bc::Base.Broadcast.Broadcasted) = copyto!(similar(bc, _source_eltype(bc)), bc) # The type of a reduction's first partial value when no `init` is given. -_first_type(op, ::Type{M}) where {M} = Base.promote_op(Base.reduce_first, Core.Typeof(op), M) +_first_type(op::OP, ::Type{M}) where {OP, M} = + Base.promote_op(Base.reduce_first, Core.Typeof(op), M) # The accumulator type of a reduction of mapped elements of type `M` into `S` (the type of `init` # or of the destination, `Union{}` for neither): the types Base's fold `op(op(S, M), M)...` goes @@ -30,7 +31,7 @@ _first_type(op, ::Type{M}) where {M} = Base.promote_op(Base.reduce_first, Core.T # `Int`, as Base's `add_sum` does, and partial results have one type however many elements they # combine. `Union{}` when inference shows that the fold always throws; only a reduction with # elements to combine is an error then (`_check_acctype`). -function _reduce_acctype(op, ::Type{S}, ::Type{M}) where {S, M} +function _reduce_acctype(op::OP, ::Type{S}, ::Type{M}) where {OP, S, M} A = _first_type(op, M) if S !== Union{} B = Base.promote_op(op, S, M) @@ -58,8 +59,8 @@ _check_acctype(op, f, ::Type{A}) where {A} = A === Union{} && throw(ArgumentErro # a partial result takes (`Base.reduce_first`'s of an element, or `op`'s of a partial result and # an element or of two partial results) has no conversion to it at all; whether the values fit is # the caller's obligation. -_acctype(op, ::Type{S}, ::Type{M}, ::Nothing) where {S, M} = _reduce_acctype(op, S, M) -function _acctype(op, ::Type{S}, ::Type{M}, ::Type{A}) where {S, M, A} +_acctype(op::OP, ::Type{S}, ::Type{M}, ::Nothing) where {OP, S, M} = _reduce_acctype(op, S, M) +function _acctype(op::OP, ::Type{S}, ::Type{M}, ::Type{A}) where {OP, S, M, A} M === Union{} && return A # (no elements to hold) _check_holds(A, _first_type(op, M), "a one-element partial result (`Base.reduce_first`)") for (T, what) in ((Base.promote_op(op, A, M), "a partial result and an element"), @@ -161,7 +162,7 @@ end # GPUArraysCore's neutral element for `op` when it has one, else an empty lane. A seed that is not # a lane stands for the accumulator type, so an abstract one (only on the host) is a lane: holding # the caller's `neutral`, or empty. -function _reduce_seed(op, ::Type{A}, neutral) where {A} +function _reduce_seed(op::OP, ::Type{A}, neutral) where {OP, A} isconcretetype(A) || return neutral === nothing ? _Lane{A}() : _Lane{A}(neutral) neutral === nothing || return convert(A, neutral) Base.promote_op(neutral_element, Core.Typeof(op), Type{A}) === Union{} && return _Lane{A}() @@ -184,6 +185,10 @@ 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 +struct _NoFinish end +@inline _finish(op, ::_NoFinish, dst, i, partial) = partial + # Unrolled map constructing a tuple @inline function unrolled_map_index(f, tuple_vector::Tuple) @@ -204,7 +209,8 @@ end # Reductions whose every output reduces a single element, a partial result of the accumulator # type `A`: `dst` and `src` have the same length. -function _mapreduce_nd_single!(f, op, dst, src, backend, ::Type{A}; init, launch...) where {A} +function _mapreduce_nd_single!(f::F, op::OP, dst, src, backend, ::Type{A}; init, + launch...) where {F, OP, A} _foreachindex(eachindex(dst), backend; launch...) do i x = Base.mapreduce_first(f, op, src[i]) dst[i] = _finish(op, init, dst, i, A === Union{} ? x : convert(A, x)) diff --git a/test/generic/reduce.jl b/test/generic/reduce.jl index 78b0f507..c0d1e464 100644 --- a/test/generic/reduce.jl +++ b/test/generic/reduce.jl @@ -1050,12 +1050,13 @@ end @testset "reduce block sizes" begin # The block-level tree reduction must cover every power-of-two group size the device allows; # each shape below reaches a different kernel: the 1D block reduction, the grid-strided - # one-block-per-output reduction (contiguous and strided sources), and the multi-block - # reduction with its second pass. + # one-block-per-output reduction (contiguous and strided sources), the multi-block + # reduction with its second pass, and the column reduction with its second pass. Random.seed!(0) vh = rand(Int32(1):Int32(100), 100_000) v = array_from_host(vh) - shapes = (((3000, 40), 1), ((40, 3000), 2), ((20_000, 4), 1)) + shapes = (((3000, 40), 1), ((40, 3000), 2), ((20_000, 4), 1), ((64, 8192), 2), + ((10, 30_000), 2), ((3000, 4000), 2)) mats = [(array_from_host(rand(Int32(1):Int32(100), sz)), dims) for (sz, dims) in shapes] block_sizes = TEST_KERNELS ? filter(<=(MAX_BLOCK_SIZE), 2 .^ (0:10)) : [nothing] for block_size in block_sizes @@ -1065,6 +1066,54 @@ end @test Array(AK.reduce(+, m; init=Int32(0), dims, alg)) == sum(Array(m); dims) end end + # The column reduction in one pass, which only more outputs than a test affords choose + if TEST_KERNELS + h = rand(Int32(1):Int32(100), 64, 1000) + R = array_from_host(fill(Int32(1), 64)) + AK._mapreduce_nd_launch!( + :columns, 1, identity, +, R, array_from_host(h), BACKEND, 256, 256; + init=AK._Fold(), neutral=Int32(0), bufs=(;), base_offset=0, reduce_strides=(64,), + reduce_sizes=(1000,), outer_strides=(1,), outer_sizes=(64,), dst_size=64, + reduce_size=1000) + @test Array(R) == vec(sum(h; dims=2)) .+ 1 + end +end + + +@testset "column reductions" begin + # Few contiguous outputs, each reducing a long strided row, take the column reduction in two + # passes (with enough elements for every backend's threshold): a tile of 8 outputs, and two + # tiles of 32 outputs the second of which is partial + Random.seed!(0) + alg = REDUCE_ALG + for (m, n) in ((8, 600_000), (40, 110_000)) + h = rand(Int8(-9):Int8(9), m, n) + a = array_from_host(h) + s = sum(Int.(h); dims=2) + if TEST_KERNELS + segs = AK._mapreduce_segments((m, n), (1, m), (2,)) + @test AK._mapreduce_nd_shape(m, n, segs, 256, AK.reduce_tuning(BACKEND, Int))[1] === + :columns + end + # `init` applied once, folding into `R`, into an `R` narrower than the accumulator + R = array_from_host(fill(Int16(1), m)) + AK.mapreducedim!(identity, +, R, a; acctype=Int32, alg) + @test Array(R) == vec(s) .+ 1 + AK.mapreducedim!(identity, +, R, a; acctype=Int32, init=Int16(7), alg) + @test Array(R) == vec(s) .+ 7 + # ... with a workspace, used twice + ws = AK.workspace(AK.mapreducedim!, identity, +, R, a; acctype=Int32, alg) + AK.mapreducedim!(identity, +, R, a; acctype=Int32, overwrite=true, alg, workspace=ws) + @test Array(R) == vec(s) + AK.mapreducedim!(identity, +, R, a; acctype=Int32, alg, workspace=ws) + @test Array(R) == 2 .* vec(s) + # Without a neutral element + @test Array(AK.reduce((x, y) -> max(x, y), a; dims=2, alg)) == maximum(h; dims=2) + # A strided view at an offset + p = array_from_host(rand(Int8(-9):Int8(9), m + 2, n + 3)) + v = view(p, 2:m + 1, 4:n + 3) + @test Array(AK.sum(v; dims=2, alg)) == sum(Int.(Array(p)[2:m + 1, 4:n + 3]); dims=2) + end end