diff --git a/benchmark/mapreduce_nd.jl b/benchmark/mapreduce_nd.jl index 26578be6..dac34175 100644 --- a/benchmark/mapreduce_nd.jl +++ b/benchmark/mapreduce_nd.jl @@ -22,3 +22,12 @@ for T in [UInt32, Int64, Float32] _group["base_dims=2_sin"] = @benchmarkable @sb(Base.mapreduce(sin, +, v; init=$T(0), dims=2)) setup=(v = ArrayType(rand(rng, $randrange, n1, n2))) _group["acck_dims=2_sin"] = @benchmarkable @sb(AK.mapreduce(sin, +, v; init=$T(0), dims=2)) setup=(v = ArrayType(rand(rng, $randrange, n1, n2))) end + +# A fused two-input source, which has no dense buffer: into one output, and into few outputs of +# many elements each +let _group = addgroup!(group, "fused"), n = 128 + for dims in ((1, 2, 3), (1, 2)) + _group["base_dims=$dims"] = @benchmarkable @sb(Base.mapreduce(*, +, v, w; dims=$dims)) setup=(v = ArrayType(rand(rng, Float32, $n, $n, $n)); w = ArrayType(rand(rng, Float32, $n, $n, $n))) + _group["acck_dims=$dims"] = @benchmarkable @sb(AK.mapreduce(*, +, v, w; dims=$dims)) setup=(v = ArrayType(rand(rng, Float32, $n, $n, $n)); w = ArrayType(rand(rng, Float32, $n, $n, $n))) + end +end diff --git a/src/reduce/mapreduce_1d_gpu.jl b/src/reduce/mapreduce_1d_gpu.jl index 66d4c9f3..b55fba69 100644 --- a/src/reduce/mapreduce_1d_gpu.jl +++ b/src/reduce/mapreduce_1d_gpu.jl @@ -136,7 +136,8 @@ _on_host(x::AbstractArray) = Array(x) _on_host(x::AbstractRange) = x _on_host(x) = x -_mapreduce_1d_src_view(src::AbstractArray) = @view src[1:end] +# (a view of all elements, indexed from 1 also when the source's indices are not) +_mapreduce_1d_src_view(src::AbstractArray) = @view src[begin:end] _mapreduce_1d_src_view(src::Base.Broadcast.Broadcasted) = src # The number of partial results of `mapreduce_1d_gpu` over `len` elements: one per block of the diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index 2c8fcb20..cba951e8 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -186,12 +186,13 @@ end _columns_cols(dst_size) = min(nextpow(2, dst_size), COLUMNS_COLS) # The source of a reduction along `dims` into a single output as a linearly indexed collection -# of its elements, for the whole-array kernels (`mapreduce_1d_gpu`): a `Broadcasted` object, or a -# contiguous range of a dense buffer; `nothing` for other sources. +# of its elements, for the whole-array kernels (`mapreduce_1d_gpu`): a `Broadcasted` object, an +# array without a dense buffer, or a contiguous range of a dense buffer; `nothing` for other +# sources. _mapreduce_nd_linear(src::Base.Broadcast.Broadcasted, dims_valid) = src function _mapreduce_nd_linear(src::AbstractArray, dims_valid) layout = _mapreduce_strided_layout(src) - isnothing(layout) && return nothing + isnothing(layout) && return src buffer, base_offset, src_strides = layout (rstrides, _, nr), _ = _mapreduce_segments(size(src), src_strides, dims_valid) nr == 1 && rstrides[1] == 1 || return nothing @@ -212,9 +213,7 @@ function _mapreduce_nd_sizes(src, backend, alg, ::Type{A}, neutral, dims_valid, reduce_tuning(backend, A).target_blocks, 0; device=true) return n == 0 ? (;) : (; partials=_buffer(typeof(neutral), n)) end - layout = _mapreduce_strided_layout(src) - isnothing(layout) && return (;) - segs = _mapreduce_segments(src_sizes, layout[3], dims_valid) + segs = _mapreduce_segments(src_sizes, _mapreduce_layout(src)[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 (;) @@ -300,22 +299,26 @@ function mapreduce_nd!( if dst_size == 1 linear = _mapreduce_nd_linear(src, dims_valid) if linear !== nothing - mapreduce_1d_gpu(f, op, linear, backend; init, neutral, block_size, - items_per_thread=alg.items_per_thread, max_blocks=target_blocks, - partials=get(bufs, :partials, nothing), switch_below=0, dst) + Base.inferencebarrier(mapreduce_1d_gpu)( + f, op, linear, backend; init, neutral, block_size, + items_per_thread=alg.items_per_thread, max_blocks=target_blocks, + partials=get(bufs, :partials, nothing), switch_below=0, dst) return dst end end - # The stride-based fast paths below index a flat buffer at - # `buffer[base_offset + Σ coordᵈ·strideᵈ + 1]`. This works for any source backed - # by a single dense column-major buffer (dense arrays, but also strided views, - # adjoints, permuted dims, and reshapes over one). Sources without such a buffer - # — `Broadcasted`, lazy/computed arrays — take the generic Cartesian-indexed - # fallback, which makes no layout assumption (and crucially does not wrap the - # source in `@Const`, which is what makes e.g. PermutedDimsArray uncompilable). + # The kernels below index a flat buffer at `buffer[base_offset + Σ coordᵈ·strideᵈ + 1]` + # (see `_mapreduce_layout`) layout = _mapreduce_strided_layout(src) - if isnothing(layout) + buffer, base_offset, src_strides = + isnothing(layout) ? _mapreduce_linear_layout(src) : 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, tuning) + + # One output per thread of a source without a dense buffer: iterating over the Cartesian + # indices of each output's elements saves decoding every linear index + if isnothing(layout) && shape === :by_thread blocks = cld(dst_size, block_size) kernel! = _mapreduce_nd_generic!(backend, block_size) Base.inferencebarrier(kernel!)( @@ -326,15 +329,11 @@ function mapreduce_nd!( return dst end - 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, tuning) _with_segments(segs) do reduce_strides, reduce_sizes, outer_strides, outer_sizes Base.inferencebarrier(_mapreduce_nd_launch!)( Val(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, + outer_sizes, dst_size, reduce_size, dense_buffer=!isnothing(layout), ) end @@ -348,10 +347,13 @@ end # and without the barrier a reduction with a new `f`, `op` or element type infers every launch it # could make: all shapes, segment layouts and vector widths. A dynamic dispatch costs a fraction of # a microsecond. +# +# `dense_buffer` is whether `buffer` is the dense buffer of the source, which the by-block kernel +# can read with vector loads, or a linearly indexed lazy source (see `_mapreduce_layout`). function _mapreduce_nd_launch!( ::Val{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, + dst_size, reduce_size, dense_buffer, ) where {shape, F, OP} if shape === :columns groups = reduce_groups @@ -427,7 +429,7 @@ function _mapreduce_nd_launch!( _launch_mapreduce_nd_by_block!( backend, block_size, buffer, dst, f, op, neutral, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, - dst_size, reduce_size, launch_blocks, + dst_size, reduce_size, launch_blocks, dense_buffer, ) end @@ -440,13 +442,28 @@ _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))...))...) +# The layout the reduction kernels index `src` with, as `(buffer, base_offset, strides)`: the +# strided layout of a source backed by a dense buffer, else (for a lazy array or a `Broadcasted` +# object) the source indexed linearly, as a dense array of its size. Indexing a lazy source +# linearly costs a decode of the index per element, but it gets the launch shapes of a dense +# array, which can spread a reduction into few outputs across many threads. +function _mapreduce_layout(src) + layout = _mapreduce_strided_layout(src) + return isnothing(layout) ? _mapreduce_linear_layout(src) : layout +end +_mapreduce_linear_layout(src) = + (_mapreduce_linear_source(src), 0, _mapreduce_dense_strides(size(src))) +# (a reshaped array decodes a linear index with multiplicative inverses rather than divisions) +_mapreduce_linear_source(src::AbstractArray) = vec(src) +_mapreduce_linear_source(src::Base.Broadcast.Broadcasted) = src + # 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 # source backed by a single dense column-major buffer — a dense array, but also a # strided view, adjoint, permuted-dims, or reshape over one. Returns -# `(buffer, base_offset, strides)` for such sources, or `nothing` (→ generic -# Cartesian-indexed fallback) for `Broadcasted` and anything not backed by a dense -# buffer (lazy/computed arrays, complex adjoints without `strides`, nested wrappers). +# `(buffer, base_offset, strides)` for such sources, or `nothing` (→ linear indexing, see +# `_mapreduce_layout`) for `Broadcasted` and anything not backed by a dense buffer +# (lazy/computed arrays, complex adjoints without `strides`, nested wrappers). function _mapreduce_strided_layout(src::AbstractArray) s = try strides(src) @@ -530,11 +547,11 @@ end function _launch_mapreduce_nd_by_block!( 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, + output_size, reduce_size, num_blocks, dense_buffer, ) where {F, OP} - W = _contiguous_vector_width( + W = dense_buffer ? _contiguous_vector_width( buffer, base_offset, outer_strides, reduce_strides, reduce_size, block_size, - ) + ) : 0 if W == 4 kernel! = _mapreduce_nd_by_block_contiguous!(backend, block_size) Base.inferencebarrier(kernel!)( @@ -645,14 +662,9 @@ end off end -# GPU kernel: generic fallback — one thread per output element, reducing sequentially -# over the reduced extents using Cartesian indexing. Makes no assumption about the -# source's memory layout, so it handles strided views, adjoints, permuted dims, and -# broadcasts over them. Mirrors the CPU `_mapreduce_nd_cpu_sections!` path. -# -# NOTE: the source is deliberately NOT marked `@Const` here. `@Const` prevents the -# `@inbounds` getindex of some wrappers (e.g. PermutedDimsArray's `genperm`) from -# being elided, leaving a `throw` that the GPU backends cannot compile. +# GPU kernel: one thread per output element, reducing sequentially over the reduced extents +# using Cartesian indexing; the by_thread kernel of sources without a dense buffer. Mirrors the +# CPU `_mapreduce_nd_cpu_sections!` path. @kernel inbounds=true cpu=false unsafe_indices=true function _mapreduce_nd_generic!( src, dst, f, op, neutral, init, diff --git a/test/Project.toml b/test/Project.toml index 71c42482..47157986 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,5 +1,6 @@ [deps] AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1" +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -20,5 +21,6 @@ pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" AcceleratedKernels = {path = ".."} [compat] +Adapt = "4" ParallelTestRunner = "2" oneAPI = "2.9.2" diff --git a/test/generic/reduce.jl b/test/generic/reduce.jl index 121e5402..e980cbd0 100644 --- a/test/generic/reduce.jl +++ b/test/generic/reduce.jl @@ -14,6 +14,29 @@ end # Only for backend-agnostic initialisation with KernelAbstractions.zero Base.zero(::Type{Point}) = Point(0.0f0, 0.0f0) +using Adapt: Adapt, adapt + +# A lazy array, without a dense buffer: its elements are computed from two arrays +struct LazyProduct{T, N, A <: AbstractArray{T, N}} <: AbstractArray{T, N} + a::A + b::A +end +Base.size(p::LazyProduct) = size(p.a) +Base.@propagate_inbounds Base.getindex(p::LazyProduct{T, N}, I::Vararg{Int, N}) where {T, N} = + p.a[I...] * p.b[I...] +Adapt.adapt_structure(to, p::LazyProduct) = LazyProduct(adapt(to, p.a), adapt(to, p.b)) +KernelAbstractions.get_backend(p::LazyProduct) = get_backend(p.a) + +# A lazy vector whose indices start at 2 +struct ShiftedVector{T, A <: AbstractVector{T}} <: AbstractVector{T} + a::A +end +Base.axes(v::ShiftedVector) = (2:length(v.a) + 1,) +Base.size(v::ShiftedVector) = size(v.a) +Base.@propagate_inbounds Base.getindex(v::ShiftedVector, i::Int) = v.a[i - 1] +Adapt.adapt_structure(to, v::ShiftedVector) = ShiftedVector(adapt(to, v.a)) +KernelAbstractions.get_backend(v::ShiftedVector) = get_backend(v.a) + @testset "reduce_1d" begin Random.seed!(0) @@ -352,7 +375,7 @@ end else # Strided GPU sources (views, adjoints, permuted dims) take the stride-based # fast path over their dense parent buffer; the offset view exercises a nonzero - # base offset. Broadcasted/lazy sources still take the generic fallback. + # base offset. vh = reshape(Int32(1):Int32(40), 5, 8) v = array_from_host(vh) @test Array(AK.reduce(+, @view(v[:, 1:2:end]); alg=REDUCE_ALG, init=Int32(0), dims=2)) == @@ -773,7 +796,7 @@ end else # Strided GPU sources (views, adjoints, permuted dims) take the stride-based # fast path over their dense parent buffer; the offset view exercises a nonzero - # base offset. Broadcasted/lazy sources still take the generic fallback. + # base offset. vh = reshape(Int32(1):Int32(40), 5, 8) v = array_from_host(vh) @test Array(AK.mapreduce(x -> x - Int32(1), +, @view(v[:, 1:2:end]); alg=REDUCE_ALG, init=Int32(0), dims=2)) == @@ -1074,7 +1097,7 @@ end Val(: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) + reduce_size=1000, dense_buffer=true) @test Array(R) == vec(sum(h; dims=2)) .+ 1 end end @@ -1165,6 +1188,52 @@ end end +@testset "lazy sources" begin + # Sources without a dense buffer, a lazy array and a fused one, reduced into one output and + # with every launch shape along `dims`: few outputs of many elements each (in one block per + # output, several, or one block for a few), contiguous outputs of strided elements (in + # columns, or tiles), non-adjacent dimensions, and many outputs of few elements each + Random.seed!(0) + alg = REDUCE_ALG + for (shape, dims) in (((64, 64, 64), (1, 2, 3)), ((64, 64, 64), (1, 2)), + ((300, 300, 3), (1, 2)), ((200, 5), 1), ((64, 64, 64), (2, 3)), + ((256, 256), 2), ((64, 64, 64), (1, 3)), ((64, 64, 64), 3)) + ha, hb = rand(Int32(-9):Int32(9), shape...), rand(Int32(-9):Int32(9), shape...) + da, db = array_from_host(ha), array_from_host(hb) + lazy, h = LazyProduct(da, db), ha .* hb + @test Array(AK.reduce(+, lazy; dims, alg)) == sum(h; dims) + @test Array(AK.reduce(+, lazy; dims, init=Int32(7), alg)) == reduce(+, h; dims, init=7) + # (without a neutral element) + @test Array(AK.reduce((a, b) -> max(a, b), lazy; dims, alg)) == maximum(h; dims) + @test Array(AK.mapreduce(*, +, da, db; dims, alg)) == sum(h; dims) + # Folding into the destination's values, and overwriting them + hR = rand(Int32(-9):Int32(9), size(sum(h; dims))) + R = array_from_host(hR) + AK.mapreducedim!(identity, +, R, lazy; alg) + @test Array(R) == hR .+ sum(h; dims) + AK.mapreducedim!(abs, max, R, lazy; overwrite=true, alg) + @test Array(R) == maximum(abs, h; dims) + end + + # ... whose indices do not start at 1 + h = rand(Int32(-9):Int32(9), 10_000) + shifted = ShiftedVector(array_from_host(h)) + @test AK.reduce(+, shifted; alg) == sum(h) + @test Array(AK.reduce(+, shifted; dims=1, alg)) == [sum(h)] + + # Reductions into one or few outputs split each output's reduction across blocks, and keep + # their partial results + if TEST_KERNELS + lazy = LazyProduct(array_from_host(ones(Int32, 300, 300, 3)), + array_from_host(ones(Int32, 300, 300, 3))) + for R in (array_from_host(zeros(Int32, 1, 1, 1)), array_from_host(zeros(Int32, 1, 1, 3))) + @test haskey(AK.workspace_size(AK.mapreducedim!, identity, +, R, lazy; + alg=AK.BlockReduce()), :partials) + end + end +end + + @testset "reduction contract" begin Random.seed!(0) alg = REDUCE_ALG