Reduce lazy sources along dims with the launch shapes of dense ones

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

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

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,

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"

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