From bd90968d6d9f168ea04a5c3cb85b9fb13d323112 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 8 Oct 2026 19:16:21 +0200 Subject: [PATCH 1/3] Reduce lazy sources along dims with the launch shapes of dense ones A source without a dense buffer, such as a lazy array or a Broadcasted object, was reduced along dims by a kernel with one thread per output, whatever the number of outputs. Reducing such a source into one element therefore ran on a single thread: a dot product of two lazy 128^3 fields, as Oceananigans computes them, took 220 ms instead of 23 us. Index these sources linearly, as a dense array of their size, so that they take the launch shapes of a dense array: the whole-array kernels for one output, and the cooperative kernels for few outputs. Only the shape with one output per thread keeps the Cartesian-indexed kernel, which needs no index decoding. A lazy array is indexed through a reshape, which decodes linear indices with multiplicative inverses, and is never read with the by-block kernel's vector loads. This also fixes whole-array reductions (`dims=:`) of vectors whose indices do not start at 1, which viewed their elements from index 1. --- src/reduce/mapreduce_1d_gpu.jl | 3 +- src/reduce/mapreduce_nd.jl | 81 +++++++++++++++++++--------------- test/Project.toml | 2 + test/generic/reduce.jl | 75 +++++++++++++++++++++++++++++-- 4 files changed, 122 insertions(+), 39 deletions(-) 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..62e8f334 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 (;) @@ -307,15 +306,18 @@ function mapreduce_nd!( 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 +328,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 +346,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 +428,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 +441,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 +546,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 +661,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 From 280a90ae2888188149259a5e321a48e5aa27db60 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 8 Oct 2026 19:16:26 +0200 Subject: [PATCH 2/3] Infer the whole-array launch of a reduction along dims only when taken A reduction along dims into one element calls the whole-array kernels, whose host code was inferred for every new `f`, `op` and element type, whether the reduction took that path or not. Call it through `Base.inferencebarrier`, like the other launches of `mapreduce_nd!`. This saves about 0.05 s of the first call of a reduction along dims with a new `f` on CUDA, and makes up for the inference that lazy sources now share with dense ones. --- src/reduce/mapreduce_nd.jl | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index 62e8f334..cba951e8 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -299,9 +299,10 @@ 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 From 8a348156437fde4890555830162fd8ca26af965f Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 8 Oct 2026 19:16:26 +0200 Subject: [PATCH 3/3] Benchmark reductions of a fused source along dims A two-input reduction reads a fused Broadcasted source, which has no dense buffer. Measure it into one output and into few outputs of many elements each, the cases that used to run one thread per output. --- benchmark/mapreduce_nd.jl | 9 +++++++++ 1 file changed, 9 insertions(+) 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