From 725b64ca38452656731edd7f1a2dccd6f9f8000a Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 15:52:32 +0200 Subject: [PATCH] Fix sparse mapreduce for operators other than + MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The row and column reduction kernels started from op(init, init) with init = zero(eltype), so maximum of a row whose stored entries are all negative returned zero, and they folded the implicit zeros in as nzeros * f(0), which is only right for +: mapreduce(x -> x + 1, max, A) counted the implicit zeros instead of taking their maximum. Reduce the stored values of each slice starting from the first one, and fold in the implicit zeros afterwards with op itself (by repeated squaring, or once for idempotent operators), as SparseArrays does. The full reduction now reduces the stored values and folds the zeros in on the host instead of reducing per-row results with a zero init. An explicit init is combined with each result, empty slices follow Base, and dims=2 returns an m×1 matrix like Base instead of a vector. --- src/host/sparse.jl | 180 +++++++++++++++++++-------------------- test/testsuite/sparse.jl | 59 +++++++++++++ 2 files changed, 148 insertions(+), 91 deletions(-) diff --git a/src/host/sparse.jl b/src/host/sparse.jl index 56c3599b0..be3133001 100644 --- a/src/host/sparse.jl +++ b/src/host/sparse.jl @@ -925,71 +925,65 @@ function Broadcast.copy(bc::Broadcasted{<:Union{GPUSparseVecStyle,GPUSparseMatSt kernel(args...; ndrange) return output end -## COV_EXCL_START -@kernel function csr_reduce_kernel(f::F, op::OP, neutral, zeros_preserved::Bool, output::DenseArray, args...) where {F, OP} - # every thread processes an entire row - row = @index(Global, Linear) - if row ≤ size(output, 1) - iter = @inbounds CSRIterator{Int}(row, args...) - - val = op(neutral, neutral) +# `op` folded over `n ≥ 1` copies of `x`, as the implicit zeros of a sparse array contribute +# to a reduction: by repeated squaring for an associative `op` (e.g. `n * x` for `+`), and a +# single `x` for idempotent ones. +@inline function _fold_repeated(op, x, n::Integer) + _isidempotent(op) && return x + result = x + power = x + n -= one(n) + while n > 0 + if isodd(n) + result = op(result, power) + end + power = op(power, power) + n >>= 1 + end + return result +end +_isidempotent(op) = false +_isidempotent(::Union{typeof(max), typeof(min), typeof(|), typeof(&)}) = true - # reduce the values for this row - for (col, ptrs) in iter - I = CartesianIndex(row, col) - vals = ntuple(Val(length(args))) do i - arg = @inbounds args[i] - ptr = @inbounds ptrs[i] - _getindex(arg, I, ptr) - end - val = op(val, f(vals...)) +## COV_EXCL_START +# reduce the stored values `nzVal[first_ptr:last_ptr]` of a slice that also contains +# `nzeros` implicit zeros, and has at least one element. As in SparseArrays, the stored +# values are reduced first and the implicit zeros folded in afterwards, so `op` is assumed +# to be commutative as well as associative. +@inline function _reduce_slice(f, op, ::Type{T}, nzVal, first_ptr, last_ptr, nzeros) where {T} + if first_ptr <= last_ptr + val = convert(T, f(@inbounds nzVal[first_ptr])) + for ptr in first_ptr+one(first_ptr):last_ptr + val = convert(T, op(val, f(@inbounds nzVal[ptr]))) end - if !zeros_preserved - f_zero_val = f(zero(neutral)) - next_row_ind = row+1 - nzs_this_row = ntuple(Val(length(args))) do i - max_n_zeros = size(args[i], 2) - arg_row_ptr = args[i].rowPtr - nz_this_row = max_n_zeros - (@inbounds(arg_row_ptr[next_row_ind]) - @inbounds(arg_row_ptr[row])) - nz_this_row * f_zero_val - end - val = op(val, nzs_this_row...) + if nzeros > 0 + val = convert(T, op(val, _fold_repeated(op, convert(T, f(zero(eltype(nzVal)))), nzeros))) end + val + else + _fold_repeated(op, convert(T, f(zero(eltype(nzVal)))), nzeros) + end +end - @inbounds output[row] = val +@kernel function csr_reduce_kernel(f::F, op::OP, output::AbstractArray{T}, A) where {F, OP, T} + # every thread processes an entire row + row = @index(Global, Linear) + if row ≤ size(A, 1) + first_ptr = @inbounds A.rowPtr[row] + last_ptr = @inbounds A.rowPtr[row+1] - one(first_ptr) + nzeros = size(A, 2) - (last_ptr - first_ptr + 1) + @inbounds output[row] = _reduce_slice(f, op, T, A.nzVal, first_ptr, last_ptr, nzeros) end end -@kernel function csc_reduce_kernel(f::F, op::OP, neutral, zeros_preserved::Bool, output::DenseArray, args...) where {F, OP} +@kernel function csc_reduce_kernel(f::F, op::OP, output::AbstractArray{T}, A) where {F, OP, T} # every thread processes an entire column col = @index(Global, Linear) - if col ≤ size(output, 2) - iter = @inbounds CSCIterator{Int}(col, args...) - - val = op(neutral, neutral) - - # reduce the values for this col - for (row, ptrs) in iter - I = CartesianIndex(row, col) - vals = ntuple(Val(length(args))) do i - arg = @inbounds args[i] - ptr = @inbounds ptrs[i] - _getindex(arg, I, ptr) - end - val = op(val, f(vals...)) - end - if !zeros_preserved - f_zero_val = f(zero(neutral)) - next_col_ind = col+1 - nzs_this_col = ntuple(Val(length(args))) do i - max_n_zeros = size(args[i], 1) - arg_col_ptr = args[i].colPtr - nz_this_col = max_n_zeros - (@inbounds(arg_col_ptr[next_col_ind]) - @inbounds(arg_col_ptr[col])) - nz_this_col * f_zero_val - end - val = op(val, nzs_this_col...) - end - @inbounds output[col] = val + if col ≤ size(A, 2) + first_ptr = @inbounds A.colPtr[col] + last_ptr = @inbounds A.colPtr[col+1] - one(first_ptr) + nzeros = size(A, 1) - (last_ptr - first_ptr + 1) + @inbounds output[col] = _reduce_slice(f, op, T, A.nzVal, first_ptr, last_ptr, nzeros) end end ## COV_EXCL_STOP @@ -998,49 +992,53 @@ csc_type(A::AbstractGPUSparseMatrix) = csc_type(typeof(A)) csr_type(A::AbstractGPUSparseMatrix) = csr_type(typeof(A)) # TODO: implement mapreducedim! -function Base.mapreduce(f, op, A::AbstractGPUSparseMatrix; dims=:, init=nothing) - # figure out the destination container type by looking at the initializer element, - # or by relying on inference to reason through the map and reduce functions - if init === nothing - ET = Broadcast.combine_eltypes(f, (A,)) - ET = Base.promote_op(op, ET, ET) - (ET === Union{} || ET === Any) && - error("mapreduce cannot figure the output element type, please pass an explicit init value") - - init = zero(ET) - else - ET = typeof(init) +function Base.mapreduce(f, op, A::AbstractGPUSparseMatrix{Tv}; dims=:, init=nothing) where {Tv} + in(dims, [Colon(), 1, 2]) || error("only dims=:, dims=1 or dims=2 is supported") + + # every slice reduction `r` becomes `op(init, r)`; empty slices give `init`, or the + # reduction of an empty collection as in Base + empty_result() = init === nothing ? Base.mapreduce_empty(f, op, Tv) : init + + # figure out the destination element type by relying on inference to reason through + # the map and reduce functions, or by looking at the initializer + ET = Broadcast.combine_eltypes(f, (A,)) + ET = Base.promote_op(op, ET, ET) + if init !== nothing + ET = Base.promote_op(op, typeof(init), ET) end + (ET === Union{} || ET === Any) && + error("mapreduce cannot figure the output element type, please pass an explicit init value") - f_preserves_zeros = ( f(zero(ET)) == zero(ET) ) - # we only handle reducing along one of the two dimensions, - # or a complete reduction (requiring an additional pass) - in(dims, [Colon(), 1, 2]) || error("only dims=:, dims=1 or dims=2 is supported") + if dims == Colon() + # reduce the stored values, then fold in the implicit zeros + length(A) == 0 && return empty_result() + nzeros = length(A) - nnz(A) + zeros = nzeros > 0 ? _fold_repeated(op, convert(ET, f(zero(Tv))), nzeros) : nothing + if nnz(A) > 0 + stored = init === nothing ? mapreduce(f, op, nonzeros(A)) : + mapreduce(f, op, nonzeros(A); init) + return zeros === nothing ? stored : op(stored, zeros) + else + return init === nothing ? zeros : op(init, zeros) + end + end if A isa AbstractGPUSparseMatrixCSR && dims == 1 A = csc_type(A)(A) elseif A isa AbstractGPUSparseMatrixCSC && dims == 2 A = csr_type(A)(A) end - m, n = size(A) - val_array = nonzeros(A) - backend = get_backend(A) - output_dim = 0 - ndrange = 0 - if A isa AbstractGPUSparseMatrixCSR - output_dim = m - ndrange = m - kernel = csr_reduce_kernel(backend) - elseif A isa AbstractGPUSparseMatrixCSC - output_dim = (1, n) - ndrange = n - kernel = csc_reduce_kernel(backend) + m, n = size(A) + output = similar(nonzeros(A), ET, dims == 1 ? (1, n) : (m, 1)) + isempty(output) && return output + if size(A, dims) == 0 + return fill!(output, empty_result()) end - output = similar(val_array, ET, output_dim) - kernel(f, op, init, f_preserves_zeros, output, A; ndrange = ndrange) - if dims == Colon() - return mapreduce(identity, op, output; init) - else - return output + kernel = A isa AbstractGPUSparseMatrixCSR ? csr_reduce_kernel(get_backend(A)) : + csc_reduce_kernel(get_backend(A)) + kernel(f, op, output, A; ndrange=length(output)) + if init !== nothing + output .= op.(init, output) end + return output end diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl index e44733f3e..99b710ca8 100644 --- a/test/testsuite/sparse.jl +++ b/test/testsuite/sparse.jl @@ -310,6 +310,10 @@ function broadcasting_matrix(AT, eltypes) end end +# the result of a reduction on the host, whether it is a scalar or an array +reduced(x::Number) = x +reduced(x::AbstractArray) = Array(x) + function mapreduce_matrix(AT, eltypes) dense_AT = GPUArrays.dense_array_type(AT) @testset "SparseMatrix mapreduce" begin @@ -359,6 +363,61 @@ function mapreduce_matrix(AT, eltypes) y = mapreduce(abs, max, x) dy = mapreduce(abs, max, dx) @test y ≈ dy + + # result shapes follow Base + @test size(sum(dx; dims=1)) == (1, n) + @test size(sum(dx; dims=2)) == (m, 1) + + # functions that don't preserve zeros, with operators other than `+` + o = one(ET) + x = sprand_nozeros(ET, m, n, p) + dx = AT(x) + for dims in (:, 1, 2) + @test mapreduce(v -> v + o, *, x; dims) ≈ reduced(mapreduce(v -> v + o, *, dx; dims)) + if ET <: Real + @test mapreduce(v -> v + o, max, x; dims) ≈ reduced(mapreduce(v -> v + o, max, dx; dims)) + @test mapreduce(v -> v - o, min, x; dims) ≈ reduced(mapreduce(v -> v - o, min, dx; dims)) + end + @test mapreduce(v -> v + o, +, x; dims, init=zero(ET)) ≈ reduced(mapreduce(v -> v + o, +, dx; dims, init=zero(ET))) + @test mapreduce(v -> v + o, *, x; dims, init=o) ≈ reduced(mapreduce(v -> v + o, *, dx; dims, init=o)) + end + # an operator without a known neutral element, relying on `init` + @test mapreduce(identity, (a, b) -> a + b, dx; init=zero(ET)) ≈ mapreduce(identity, (a, b) -> a + b, x; init=zero(ET)) + y = sparse([1, 2], [1, 3], ET[1, 2], 3, 3) + dy = AT(y) + @test mapreduce(v -> v + o, +, dy) == mapreduce(v -> v + o, +, y) + @test mapreduce(v -> v + o, *, dy) == mapreduce(v -> v + o, *, y) + if ET <: Real + @test mapreduce(v -> v + o, max, dy) == mapreduce(v -> v + o, max, y) + end + + # the stored values of a row or column don't start from zero + if ET <: Real && !(ET <: Unsigned) + x = sparse(-rand(ET(1):ET(10), m, n)) + dx = AT(x) + @test nnz(dx) == m * n + for dims in (:, 1, 2) + @test maximum(x; dims) == reduced(maximum(dx; dims)) + @test minimum(-x; dims) == reduced(minimum(-dx; dims)) + end + # the same with some implicit zeros + x[1, :] .= 0 + dropzeros!(x) + dx = AT(x) + for dims in (:, 1, 2) + @test maximum(x; dims) == reduced(maximum(dx; dims)) + end + end + + # empty dimensions + for (k, l) in ((0, n), (m, 0), (0, 0)) + x = spzeros(ET, k, l) + dx = AT(x) + @test sum(dx) == sum(x) + @test collect(sum(dx; dims=1)) == sum(x; dims=1) + @test collect(sum(dx; dims=2)) == sum(x; dims=2) + @test collect(mapreduce(abs, +, dx; dims=1, init=zero(real(ET)))) == mapreduce(abs, +, x; dims=1, init=zero(real(ET))) + end end end end