From 3e63c3aceaf5d1605394d73c27d4d875d702b425 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 1 Oct 2026 15:37:20 +0200 Subject: [PATCH 1/2] Specialize the reductions' host code on their functions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Julia does not specialize a method on a function argument that it only passes on to other calls. The reductions' entry points passed `f` and `op` straight through, so the host code below them ran unspecialized: `Base.promote_op` ran inference at run time on every call to find the accumulator type, and every call after it was a dynamic dispatch. The dimension segments added more, as vectors of tuples of unknown length. Annotate the entry points with `where {F, OP}`, keep the segments in fixed-length padded tuples (dispatching the common layouts, one reduced and at most one kept segment, statically), and put a function barrier after the computation of the reduced dimensions in `mapreducedim!`. Host time of `mapreducedim!(identity, +, R, A; init=0f0)` for a 1000×10 `Float32` matrix (`dims=2`), RTX 5080: 13.9 µs before, 3.7 µs after; CUDA.jl's own kernel takes 2.6 µs, of which the launch is about 2 µs. That was the whole gap for small reductions (wall-clock minimum of synchronized runs, µs): | shape | dims | CUDA.jl | before | after | |-----------|------|---------|--------|-------| | 1000×10 | 1 | 11.8 | 20.0 | 8.1 | | 1000×10 | 2 | 6.8 | 19.9 | 7.5 | | 1000×1000 | 1 | 17.8 | 23.1 | 10.6 | | 1000×1000 | 2 | 23.8 | 21.7 | 9.6 | --- src/reduce/mapreduce_1d_cpu.jl | 4 +- src/reduce/mapreduce_1d_gpu.jl | 4 +- src/reduce/mapreduce_nd.jl | 157 ++++++++++++++++++--------------- src/reduce/reduce.jl | 55 +++++++----- src/reduce/utilities.jl | 16 ++-- 5 files changed, 133 insertions(+), 103 deletions(-) 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..43fe7baa 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -31,49 +31,59 @@ const TILED_STRIDED_ROWS_PER_BLOCK = 8 # 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}[] - +# 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: @@ -92,17 +102,16 @@ 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, 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,14 +121,12 @@ 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) @@ -152,12 +159,8 @@ function _mapreduce_nd_sizes(src, backend, alg, ::Type{A}, neutral, dims_valid, 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, - ) + segs = _mapreduce_segments(src_sizes, layout[3], dims_valid) + shape, reduce_groups = _mapreduce_nd_shape(dst_size, len, segs, alg.block_size, target_blocks) shape === :multigroup || return (;) return (; partials=_buffer(typeof(neutral), dst_size, reduce_groups)) end @@ -182,13 +185,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) @@ -253,14 +256,27 @@ 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, + target_blocks) + _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 +# 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 === :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 @@ -320,17 +336,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 +433,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 +488,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) 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/utilities.jl b/src/reduce/utilities.jl index d7cf0b6a..3ecb5803 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}() @@ -204,7 +205,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)) From c55c691873603c66f9b7199a29b0c35fe70570d8 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 1 Oct 2026 15:37:20 +0200 Subject: [PATCH 2/2] Add a column kernel for long strided reductions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A reduction into contiguous outputs along a longer strided dimension, such as `sum(A; dims=2)` of a wide column-major matrix, took the `by_block` shape: one block per output, its threads reading elements a row apart, so no two threads of a warp touched the same cache line. At 1000×100000 it ran at a quarter of the bandwidth, slower than CUDA.jl's kernel; with few outputs it took `multigroup`, as poorly coalesced. The column kernel gives each block up to 32 consecutive outputs (8 to 31 outputs take the next power of two) with the block's other threads as lanes along the reduction, so a warp's loads of one step are consecutive elements; each thread keeps 8 loads in flight, and when the outputs alone do not fill the device, each output's reduction is split across blocks and combined in a second pass, like `multigroup`. It runs blocks of at most 256 threads, as Metal allows its pipeline fewer than 1024. It is used from `columns_min_elements` elements, a new internal tuning value: 2^18 by default, 2^22 on oneAPI, where a second launch costs tens of microseconds; and from 8 outputs, as with 4 it lost to `multigroup` on Metal and oneAPI. `sum(A; dims=2)` of `Float32`, minimum of 50 synchronized runs (`CUDA.@elapsed`), RTX 5080, µs: | shape | CUDA.jl | before | after | |--------------|---------|--------|-------| | 300×10000 | 45.3 | 67.0 | 10.6 | | 1000×10000 | 141.0 | 171.2 | 18.3 | | 1000×100000 | 1459.1 | 1815.2 | 454.6 | | 5000×20000 | 1466.2 | 1849.9 | 460.6 | | 100×1000000 | 1705.9 | 1657.2 | 460.1 | | 32×3000000 | 1452.3 | 4374.9 | 424.7 | | 16×6000000 | 1078.0 | 761.5 | 428.4 | | 8×12000000 | 701.8 | 760.5 | 430.0 | --- ext/oneAPIExt.jl | 5 ++ src/reduce/mapreduce_nd.jl | 149 +++++++++++++++++++++++++++++++++---- src/reduce/tuning.jl | 7 +- src/reduce/utilities.jl | 4 + test/generic/reduce.jl | 55 +++++++++++++- 5 files changed, 203 insertions(+), 17 deletions(-) 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_nd.jl b/src/reduce/mapreduce_nd.jl index 43fe7baa..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,6 +31,18 @@ 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 the `(stride, size)` pairs of @@ -86,7 +100,12 @@ 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 @@ -102,7 +121,8 @@ _segment_tuple(t::Tuple, n::Int) = ntuple(i -> t[i], n) # 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, segs, 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 @@ -130,6 +150,19 @@ function _mapreduce_nd_shape(dst_size, reduce_size, segs, block_size, target_blo 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 @@ -149,8 +182,12 @@ function _mapreduce_nd_shape(dst_size, reduce_size, segs, block_size, target_blo 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) @@ -158,10 +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 segs = _mapreduce_segments(src_sizes, layout[3], dims_valid) - shape, reduce_groups = _mapreduce_nd_shape(dst_size, len, segs, alg.block_size, target_blocks) - shape === :multigroup || return (;) + 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 @@ -199,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")) @@ -258,8 +296,7 @@ function mapreduce_nd!( buffer, base_offset, src_strides = layout segs = _mapreduce_segments(src_sizes, src_strides, dims_valid) reduce_size = len - shape, reduce_groups = _mapreduce_nd_shape(dst_size, reduce_size, segs, 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; @@ -277,13 +314,32 @@ function _mapreduce_nd_launch!( init, neutral, bufs, base_offset, reduce_strides, reduce_sizes, outer_strides, outer_sizes, dst_size, reduce_size, ) where {F, OP} - if shape === :tiled_strided + 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) @@ -632,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/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 3ecb5803..8af4e074 100644 --- a/src/reduce/utilities.jl +++ b/src/reduce/utilities.jl @@ -185,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) 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