diff --git a/src/linalg.jl b/src/linalg.jl index efd664e5..5d792de0 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -558,42 +558,136 @@ end ## diagonal scaling +# An entry that is not stored scales to `op(0, d[j])` columnwise (`dim == 2`) or to +# `op(d[i], 0)` rowwise (`dim == 1`), which is not zero for an infinite or NaN `d[j]`. Such a +# column or row is stored in full before it is scaled, its new entries as stored zeros, so +# that it scales as a dense one does. `_fullscaled` returns a mask of the columns or rows +# that need it, or `nothing`. An eltype without a `zero` has only its stored values scaled. +function _fullscaled(A, d::AbstractVector, op::F, dim::Int) where {F} + z = _implicit_zero(A) + z isa UnknownZero && return nothing + scaleszero(dk) = _iszero(dim == 1 ? op(dk, z) : op(z, dk)) + all(scaleszero, d) && return nothing + nstored = _nstored(A, dim) + full = [!scaleszero(d[k]) && nstored[k] < size(A, 3 - dim) for k in axes(A, dim)] + return any(full) ? full : nothing +end +function _nstored(A::AbstractSparseMatrixCSC, dim::Int) + dim == 2 && return [length(nzrange(A, j)) for j in axes(A, 2)] + Arowval = getrowval(A) + counts = zeros(Int, size(A, 1)) + for j in axes(A, 2), p in nzrange(A, j) + counts[Arowval[p]] += 1 + end + return counts +end +function _nstored(x::AbstractSparseVector, dim::Int) + counts = zeros(Int, length(x)) + counts[nonzeroinds(x)] .= 1 + return counts +end + +# `A` with the columns or rows that `full` marks stored in full +function _storeinfull(A::AbstractSparseMatrixCSC, full::Vector{Bool}, dim::Int) + m, n = size(A) + Arowval, Anzval = getrowval(A), getnzval(A) + z = _densezero(A) + fullrows = dim == 1 ? findall(full) : collect(1:m) + instored = dim == 1 ? sum(j -> count(p -> full[Arowval[p]], nzrange(A, j)), axes(A, 2); init = 0) : + sum(j -> full[j] ? length(nzrange(A, j)) : 0, axes(A, 2); init = 0) + nz = nnz(A) + widemul(count(full), size(A, 3 - dim)) - instored + _checknnzfits(indtype(A), nz) + colptr = Vector{indtype(A)}(undef, n + 1) + rowval = Vector{indtype(A)}(undef, nz) + nzval = Vector{promote_type(eltype(A), typeof(z))}(undef, nz) + colptr[1] = q = 1 + for j in axes(A, 2) + rows = view(fullrows, 1:(dim == 1 || full[j] ? length(fullrows) : 0)) + q = _mergestored!(rowval, nzval, q, Arowval, Anzval, nzrange(A, j), rows, z) + colptr[j + 1] = q + end + return SparseMatrixCSC(m, n, colptr, rowval, nzval) +end +function _storeinfull(x::AbstractSparseVector, full::Vector{Bool}) + z = _densezero(x) + nz = nnz(x) + count(full) + nzind = Vector{indtype(x)}(undef, nz) + nzval = Vector{promote_type(eltype(x), typeof(z))}(undef, nz) + _mergestored!(nzind, nzval, 1, nonzeroinds(x), nonzeros(x), 1:nnz(x), findall(full), z) + return SparseVector(length(x), nzind, nzval) +end +# Write the indices `ind[ps]` and values `val[ps]`, merged with the sorted indices `inds` at +# which they store `z` where `ind` has none, to `outind` and `outval` from `q` on, and return +# the next `q`. +function _mergestored!(outind, outval, q::Int, ind, val, ps, inds, z) + p, k = Int(first(ps)), 1 + while p <= last(ps) || k <= length(inds) + i = p <= last(ps) ? Int(ind[p]) : typemax(Int) + r = k <= length(inds) ? inds[k] : typemax(Int) + outind[q] = min(i, r) + outval[q] = i <= r ? val[p] : z + i <= r && (p += 1) + r <= i && (k += 1) + q += 1 + end + return q +end + +# Scale `A`, with the columns or rows that `full` marks stored in full, into `C` with +# `scale!`, which is `_scalecols!` or `_scalerows!`. A fixed `C` is written only if it stores +# every entry of the result that is not zero. +function _scaleinfull!(scale!::S, C, A, d, op::F, full, dim::Int) where {S,F} + B = _storeinfull(A, full, dim) + _is_fixed(C) && return _copyto_fixed!(C, scale!(_writablecopy(C), B, d, op)) + return scale!(C, B, d, op) +end + # Scale the stored entries of `A` into `C`, which either is `A` or takes its pattern: # `C[i, j] = op(A[i, j], d[j])` columnwise, `C[i, j] = op(d[i], A[i, j])` rowwise. The # argument order of `op` is kept for non-commutative eltypes. The callers check that -# `length(d)` matches the scaled dimension and that `size(C) == size(A)`. +# `length(d)` matches the scaled dimension and that `size(C) == size(A)`. When a value might +# not convert to `eltype(C)`, the values go to a new vector that `C` takes once all are written. function _scalecols!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, d::AbstractVector, op::F) where {F} - copyinds!(C, A) - Cnzval = getnzval(C) + full = _fullscaled(A, d, op, 2) + isnothing(full) || return _scaleinfull!(_scalecols!, C, A, d, op, full, 2) Anzval = getnzval(A) - resize!(Cnzval, length(Anzval)) + direct = _storesafely(eltype(C), promote_op(op, eltype(A), eltype(d))) + Cnzval = direct ? _takepattern!(C, A) : Vector{eltype(C)}(undef, nnz(A)) @inbounds for col in axes(A, 2) dcol = d[col] for p in nzrange(A, col) Cnzval[p] = op(Anzval[p], dcol) end end + direct || _storescaled!(C, A, Cnzval) return C end function _scalerows!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, d::AbstractVector, op::F) where {F} - copyinds!(C, A) - Cnzval = getnzval(C) + full = _fullscaled(A, d, op, 1) + isnothing(full) || return _scaleinfull!(_scalerows!, C, A, d, op, full, 1) Anzval = getnzval(A) Arowval = getrowval(A) - resize!(Cnzval, length(Anzval)) + direct = _storesafely(eltype(C), promote_op(op, eltype(d), eltype(A))) + Cnzval = direct ? _takepattern!(C, A) : Vector{eltype(C)}(undef, nnz(A)) @inbounds for col in axes(A, 2), p in nzrange(A, col) Cnzval[p] = op(d[Arowval[p]], Anzval[p]) end + direct || _storescaled!(C, A, Cnzval) return C end function _scalerows!(x::AbstractSparseVector, d::AbstractVector, op::F) where {F} + full = _fullscaled(x, d, op, 1) + isnothing(full) || return copyto!(x, _scalerows!(_storeinfull(x, full), d, op)) xnzval = nonzeros(x) xnzind = nonzeroinds(x) + direct = _storesafely(eltype(x), promote_op(op, eltype(d), eltype(x))) + vals = direct ? xnzval : similar(xnzval) @inbounds for p in eachindex(xnzind, xnzval) - xnzval[p] = op(d[xnzind[p]], xnzval[p]) + vals[p] = op(d[xnzind[p]], xnzval[p]) end + direct || copyto!(xnzval, vals) return x end @@ -1353,6 +1447,19 @@ function copyinds!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC; copy_ end end +# Give `C` the pattern of `A`, and return the values of `C`, of the length of those of `A`. +function _takepattern!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC) + copyinds!(C, A) + return resize!(getnzval(C), length(getnzval(A))) +end +# Store `nzval` in `C` with the pattern of `A`, converting it first so that a value that does +# not convert to `eltype(C)` throws with `C` unchanged. +function _storescaled!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, nzval::AbstractVector) + vals = convert(Vector{eltype(C)}, nzval) + copyto!(_takepattern!(C, A), vals) + return C +end + """ rowcheck_index(A::AbstractSparseMatrixCSC, row::Integer, col::Integer) diff --git a/src/matmul.jl b/src/matmul.jl index 50de6da2..435a9def 100644 --- a/src/matmul.jl +++ b/src/matmul.jl @@ -348,6 +348,8 @@ function *(A::Diagonal, b::AbstractSparseVector) DimensionMismatch(lazy"The dimension of the matrix A $(size(A)) and of the vector b $(length(b))") ) end + full = _fullscaled(b, A.diag, *, 1) + isnothing(full) || return A * _storeinfull(b, full) T = promote_eltype(A, b) res = similar(b, T) nzind_b = nonzeroinds(b) @@ -421,12 +423,16 @@ end function (*)(Da::Diagonal, A::SparseMatrixCSC, Db::Diagonal) (size(Da, 2) == size(A,1) && size(A,2) == size(Db,1)) || throw(DimensionMismatch("incompatible sizes")) + da, db = map(parent, (Da, Db)) + fullrows = _fullscaled(A, da, *, 1) + isnothing(fullrows) || return Da * _storeinfull(A, fullrows, 1) * Db + fullcols = _fullscaled(A, db, *, 2) + isnothing(fullcols) || return Da * _storeinfull(A, fullcols, 2) * Db T = promote_op(matprod, eltype(Da), promote_op(matprod, eltype(A), eltype(Db))) dest = similar(A, T) vals_dest = getnzval(dest) rows = getrowval(A) vals = getnzval(A) - da, db = map(parent, (Da, Db)) @inbounds for col in axes(A,2) dbcol = db[col] for i in nzrange(A, col) @@ -828,9 +834,10 @@ Store `A*D*α + C*β`, or `D*A*α + C*β`, in the sparse matrix `C`, and return columns, or the rows, of `A` scaled by the diagonal of `D`. Unlike a dense destination, `C` takes a new sparsity pattern: that of `A` when `β` is zero, -and otherwise the union of the patterns of `A` and `C`. A fixed-pattern `C` keeps its -pattern instead, and `mul!` throws, with `C` unchanged, if a nonzero of the result falls -outside it. +and otherwise the union of the patterns of `A` and `C`. A column, or row, that an infinite or +NaN entry of `D` scales is stored in full, as its entries that `A` does not store become +`NaN`. A fixed-pattern `C` keeps its pattern instead, and `mul!` throws, with `C` unchanged, +if a nonzero of the result falls outside it. A sparse destination is supported for scaling only, by a `Diagonal` here or by a number with `mul!(C, A, b)`. Form other products of sparse matrices with `*`. With a dense @@ -844,6 +851,8 @@ function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, D::Diagona n == lb || throw(DimensionMismatch(lazy"A has size ($m, $n) but D has size ($lb, $lb)")) size(A)==size(C) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))")) iszero(alpha) && return _diagmul_scale!(C, beta) + full = _fullscaled(A, b, (a, bj) -> isone(alpha) ? a * bj : a * bj * alpha, 2) + isnothing(full) || return mul!(C, _storeinfull(A, full, 2), D, alpha, beta) R = _diagmul_eltype(C, promote_op(matprod, eltype(A), eltype(D)), alpha, beta) _storesafely(eltype(C), R) || return _diagmul_store!(C, _spdiagmul!(_diagmul_widened(C, A, R, beta), A, D, alpha, beta)) @@ -951,6 +960,8 @@ function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCS m == lb || throw(DimensionMismatch(lazy"D has size ($lb, $lb) but A has size ($m, $n)")) size(A)==size(C) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))")) iszero(alpha) && return _diagmul_scale!(C, beta) + full = _fullscaled(A, b, (bi, a) -> isone(alpha) ? bi * a : bi * a * alpha, 1) + isnothing(full) || return mul!(C, D, _storeinfull(A, full, 1), alpha, beta) R = _diagmul_eltype(C, promote_op(matprod, eltype(D), eltype(A)), alpha, beta) _storesafely(eltype(C), R) || return _diagmul_store!(C, _diagspmul!(_diagmul_widened(C, A, R, beta), D, A, alpha, beta)) @@ -1022,18 +1033,18 @@ _lmulkeepszeros(b::Number, A) = (z = _implicit_zero(A); z isa UnknownZero || _is function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, b::Number) size(A)==size(C) || throw(DimensionMismatch("A has size $(size(A)) but C has size $(size(C))")) (_is_fixed(C) || !_rmulkeepszeros(A, b)) && return broadcast!(*, C, A, b) - copyinds!(C, A) - resize!(getnzval(C), length(getnzval(A))) - mul!(getnzval(C), getnzval(A), b) + _storesafely(eltype(C), promote_op(*, eltype(A), typeof(b))) || + return _storescaled!(C, A, view(getnzval(A), 1:nnz(A)) * b) + mul!(_takepattern!(C, A), getnzval(A), b) C end function mul!(C::AbstractSparseMatrixCSC, b::Number, A::AbstractSparseMatrixCSC) size(A)==size(C) || throw(DimensionMismatch("A has size $(size(A)) but C has size $(size(C))")) (_is_fixed(C) || !_lmulkeepszeros(b, A)) && return broadcast!(*, C, b, A) - copyinds!(C, A) - resize!(getnzval(C), length(getnzval(A))) - mul!(getnzval(C), b, getnzval(A)) + _storesafely(eltype(C), promote_op(*, typeof(b), eltype(A))) || + return _storescaled!(C, A, b * view(getnzval(A), 1:nnz(A))) + mul!(_takepattern!(C, A), b, getnzval(A)) C end diff --git a/test/linalg.jl b/test/linalg.jl index d39ca267..5d569cdf 100644 --- a/test/linalg.jl +++ b/test/linalg.jl @@ -344,6 +344,15 @@ end C = spzeros(ComplexF64, 3, 3) @test_throws DimensionMismatch ldiv!(C, Dl, A) @test nnz(C) == 0 + @static if COMPREHENSIVE + # a quotient that does not convert to the eltype of the destination throws with it unchanged + A0 = sparse([2 0; 0 3]); x0 = sparsevec([1, 2], [2, 3]); D2 = Diagonal([2, 2]) + for (f!, X0) in ((C -> ldiv!(C, D2, A0), sparse([1 0; 0 0])), (B -> rdiv!(B, D2), A0), (y -> ldiv!(D2, y), x0)) + X = copy(X0) + @test_throws InexactError f!(X) + @test mismatch(X, X0) === nothing + end + end end end diff --git a/test/matmul.jl b/test/matmul.jl index 4d458f9d..984145f5 100644 --- a/test/matmul.jl +++ b/test/matmul.jl @@ -448,6 +448,12 @@ end @test mismatch(Diagonal(b) * sA, Diagonal(b) * dA) === nothing @test mismatch(mul!(sC, Diagonal(b), sA), Diagonal(b) * dA) === nothing @test mismatch(lmul!(Diagonal(b), copy(sA)), Diagonal(b) * dA) === nothing + @static if COMPREHENSIVE + # an infinite entry of the Diagonal scales the entries that are not stored to NaN, as for dense + S = sparse([1.0 0; 0 1]); Dinf = Diagonal([Inf, 1]) + @test mismatch(S * Dinf, Matrix(S) * Dinf) === nothing + @test mismatch(rmul!(copy(S), Dinf), Matrix(S) * Dinf) === nothing + end # adjoint/transpose of a sparse matrix with a Diagonal (issue #619) for (T, W) in ((ComplexF64, adjoint), (ComplexF64, transpose)) @@ -512,17 +518,18 @@ end @test mismatch(F * Dr, dA * Dr; approx=true) === nothing @test !_is_fixed(Dl * F) && !_is_fixed(F * Dr) end - # the kernels touch only the stored entries: exactly nnz(S) scalar multiplications, - # whereas the generic Diagonal kernel visits every element of the result + # the kernels touch only the stored entries: exactly nnz(S) scalar multiplications, and + # one per diagonal entry to find what it scales a zero to, whereas the generic Diagonal + # kernel visits every element of the result S = opcount_sparse(fixture(Float64, 20, 30)) Dl = Diagonal(OpCount.(fixturedense(Float64, 30))); Dr = Diagonal(OpCount.(fixturedense(Float64, 20))) for W in (adjoint,) - @test mulcount(() -> W(S) * Dr) == nnz(S) - @test mulcount(() -> Dl * W(S)) == nnz(S) + @test mulcount(() -> W(S) * Dr) == nnz(S) + 20 + @test mulcount(() -> Dl * W(S)) == nnz(S) + 30 @static if COMPREHENSIVE C = similar(W(S)) - @test mulcount(() -> mul!(C, W(S), Dr)) == nnz(S) - @test mulcount(() -> mul!(C, Dl, W(S))) == nnz(S) + @test mulcount(() -> mul!(C, W(S), Dr)) == nnz(S) + 20 + @test mulcount(() -> mul!(C, Dl, W(S))) == nnz(S) + 30 end end @static if COMPREHENSIVE @@ -573,6 +580,15 @@ end @test mismatch(mul!(sC, sA, 0.5), 0.5 * dA) === nothing @test mismatch(lmul!(0.5, copy(sA)), 0.5 * dA) === nothing @test mul!(sC, 0.5, sA) == mul!(sC, sA, 0.5) + @static if COMPREHENSIVE + # a product that does not convert to the eltype of `C` throws before `C` takes the pattern of `A` + A0 = sparse([1 0; 0 3]); C0 = sparse([2 0; 0 0]) + for f! in (C -> mul!(C, A0, 0.5), C -> mul!(C, 0.5, A0)) + C = copy(C0) + @test_throws InexactError f!(C) + @test mismatch(C, C0) === nothing + end + end @testset "inverse scaling with mul!" begin bi = inv.(b)