Speed up small reductions and reductions of wide matrices along dims

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

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

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

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,36 +116,36 @@ 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,
     acctype=nothing,
     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)
     return _mapreduce_run(f, op, _mapreduce_source(src), s, init, bufs)
 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,28 +286,34 @@ 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,
     overwrite::Bool=false,
     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;

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
 

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))

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