From aad862434a7921a5bcdfc390767a282700fab2a6 Mon Sep 17 00:00:00 2001 From: Sacha Verweij Date: Sat, 2 Jul 2016 09:27:14 -0700 Subject: [PATCH] Rewrite vectorized unary functions over SparseMatrixCSCs, leveraging higher order functions and multiple dispatch to displace eval. Fixes some apparent type instabilities. --- base/sparse/sparsematrix.jl | 216 +++++++++++++++++------------------- 1 file changed, 99 insertions(+), 117 deletions(-) diff --git a/base/sparse/sparsematrix.jl b/base/sparse/sparsematrix.jl index effbe9813afc72..edf62c4ad4d1ea 100644 --- a/base/sparse/sparsematrix.jl +++ b/base/sparse/sparsematrix.jl @@ -1033,134 +1033,116 @@ end ## Unary arithmetic and boolean operators -macro _unary_op_nz2z_z2z(op,A,Tv,Ti) - esc(quote - nfilledA = nnz($A) - colptrB = Array{$Ti}($A.n+1) - rowvalB = Array{$Ti}(nfilledA) - nzvalB = Array{$Tv}(nfilledA) - - nzvalA = $A.nzval - colptrA = $A.colptr - rowvalA = $A.rowval - - k = 0 # number of additional zeros introduced by op(A) - @inbounds for i = 1 : $A.n - colptrB[i] = colptrA[i] - k - for j = colptrA[i] : colptrA[i+1]-1 - opAj = $(op)(nzvalA[j]) - if opAj == 0 - k += 1 - else - rowvalB[j - k] = rowvalA[j] - nzvalB[j - k] = opAj - end - end - end - colptrB[end] = $A.colptr[end] - k - deleteat!(rowvalB, colptrB[end]:nfilledA) - deleteat!(nzvalB, colptrB[end]:nfilledA) - return SparseMatrixCSC($A.m, $A.n, colptrB, rowvalB, nzvalB) - end) # quote -end - -# Operations that may map nonzeros to zero, and zero to zero -# Result is sparse -for op in (:ceil, :floor, :trunc, :round, - :sin, :tan, :asin, :atan, - :sinh, :tanh, :asinh, :atanh, - :sinpi, :cosc, - :sind, :tand, :asind, :atand) - @eval begin - $(op){Tv,Ti}(A::SparseMatrixCSC{Tv,Ti}) = @_unary_op_nz2z_z2z($op,A,Tv,Ti) - end # quote -end # macro - -for op in (:real, :imag) - @eval begin - ($op){Tv<:Complex,Ti}(A::SparseMatrixCSC{Tv,Ti}) = @_unary_op_nz2z_z2z($op,A,Tv.parameters[1],Ti) - end # quote -end # macro -real{Tv<:Number,Ti}(A::SparseMatrixCSC{Tv,Ti}) = copy(A) -imag{Tv<:Number,Ti}(A::SparseMatrixCSC{Tv,Ti}) = spzeros(Tv, Ti, A.m, A.n) - -for op in (:ceil, :floor, :trunc, :round) - @eval begin - ($op){T,Tv,Ti}(::Type{T},A::SparseMatrixCSC{Tv,Ti}) = @_unary_op_nz2z_z2z($op,A,T,Ti) - end # quote -end # macro - - -# Operations that map nonzeros to nonzeros, and zeros to zeros -# Result is sparse -for op in (:-, :log1p, :expm1) - @eval begin - - function ($op)(A::SparseMatrixCSC) - B = similar(A) - nzvalB = B.nzval - nzvalA = A.nzval - @simd for i=1:length(nzvalB) - @inbounds nzvalB[i] = ($op)(nzvalA[i]) - end - return B - end - +""" +Helper macro for the unary broadcast definitions below. Takes parent method `fp` and a set +of desired child methods `fcs`, and builds an expression defining each of the child methods +such that `fc(A::SparseMatrixCSC) = fp(fc, A)`. +""" +macro _enumerate_childmethods(fp, fcs...) + fcexps = Expr(:block) + for fc in fcs + push!(fcexps.args, :( $(esc(fc))(A::SparseMatrixCSC) = $(esc(fp))($(esc(fc)), A) ) ) end + return fcexps end -function abs{Tv<:Complex,Ti}(A::SparseMatrixCSC{Tv,Ti}) - T = Tv.parameters[1] - (T <: Integer) && (T = (T <: BigInt) ? BigFloat : Float64) - @_unary_op_nz2z_z2z(abs,A,T,Ti) -end -abs2{Tv<:Complex,Ti}(A::SparseMatrixCSC{Tv,Ti}) = @_unary_op_nz2z_z2z(abs2,A,Tv.parameters[1],Ti) -for op in (:abs, :abs2) - @eval begin - function ($op){Tv<:Number,Ti}(A::SparseMatrixCSC{Tv,Ti}) - B = similar(A) - nzvalB = B.nzval - nzvalA = A.nzval - @simd for i=1:length(nzvalB) - @inbounds nzvalB[i] = ($op)(nzvalA[i]) +# Operations that map zeros to zeros and may map nonzeros to zeros, yielding a sparse matrix +""" +Takes unary function `f` that maps zeros to zeros and may map nonzeros to zeros, and returns +a new `SparseMatrixCSC{TiA,TvB}` `B` generated by applying `f` to each nonzero entry in +`A` and retaining only the resulting nonzeros. +""" +function _broadcast_unary_nz2z_z2z_T{TvA,TiA,TvB}(f::Function, A::SparseMatrixCSC{TvA,TiA}, ::Type{TvB}) + Bcolptr = Array{TiA}(A.n + 1) + Browval = Array{TiA}(nnz(A)) + Bnzval = Array{TvB}(nnz(A)) + Bk = 1 + @inbounds for j in 1:A.n + Bcolptr[j] = Bk + for Ak in nzrange(A, j) + x = f(A.nzval[Ak]) + if x != 0 + Browval[Bk] = A.rowval[Ak] + Bnzval[Bk] = x + Bk += 1 end - return B end end -end - + Bcolptr[A.n + 1] = Bk + resize!(Browval, Bk - 1) + resize!(Bnzval, Bk - 1) + return SparseMatrixCSC(A.m, A.n, Bcolptr, Browval, Bnzval) +end +function _broadcast_unary_nz2z_z2z{Tv}(f::Function, A::SparseMatrixCSC{Tv}) + _broadcast_unary_nz2z_z2z_T(f, A, Tv) +end +@_enumerate_childmethods(_broadcast_unary_nz2z_z2z, + sin, sinh, sind, asin, asinh, asind, + tan, tanh, tand, atan, atanh, atand, + sinpi, cosc, ceil, floor, trunc, round) +real(A::SparseMatrixCSC) = copy(A) +imag{Tv,Ti}(A::SparseMatrixCSC{Tv,Ti}) = spzeros(Tv, Ti, A.m, A.n) +real{TTv}(A::SparseMatrixCSC{Complex{TTv}}) = _broadcast_unary_nz2z_z2z_T(real, A, TTv) +imag{TTv}(A::SparseMatrixCSC{Complex{TTv}}) = _broadcast_unary_nz2z_z2z_T(imag, A, TTv) +ceil{To}(::Type{To}, A::SparseMatrixCSC) = _broadcast_unary_nz2z_z2z_T(ceil, A, To) +floor{To}(::Type{To}, A::SparseMatrixCSC) = _broadcast_unary_nz2z_z2z_T(floor, A, To) +trunc{To}(::Type{To}, A::SparseMatrixCSC) = _broadcast_unary_nz2z_z2z_T(trunc, A, To) +round{To}(::Type{To}, A::SparseMatrixCSC) = _broadcast_unary_nz2z_z2z_T(round, A, To) + +# Operations that map zeros to zeros and map nonzeros to nonzeros, yielding a sparse matrix +""" +Takes unary function `f` that maps zeros to zeros and nonzeros to nonzeros, and returns a +new `SparseMatrixCSC{TiA,TvB}` `B` generated by applying `f` to each nonzero entry in `A`. +""" +function _broadcast_unary_nz2nz_z2z_T{TvA,TiA,TvB}(f::Function, A::SparseMatrixCSC{TvA,TiA}, ::Type{TvB}) + Bcolptr = Vector{TiA}(A.n + 1) + Browval = Vector{TiA}(nnz(A)) + Bnzval = Vector{TvB}(nnz(A)) + copy!(Bcolptr, 1, A.colptr, 1, A.n + 1) + copy!(Browval, 1, A.rowval, 1, nnz(A)) + @inbounds @simd for k in 1:nnz(A) + Bnzval[k] = f(A.nzval[k]) + end + return SparseMatrixCSC(A.m, A.n, Bcolptr, Browval, Bnzval) +end +function _broadcast_unary_nz2nz_z2z{Tv}(f::Function, A::SparseMatrixCSC{Tv}) + _broadcast_unary_nz2nz_z2z_T(f, A, Tv) +end +@_enumerate_childmethods(_broadcast_unary_nz2nz_z2z, + log1p, expm1, abs, abs2, conj) +abs2{TTv}(A::SparseMatrixCSC{Complex{TTv}}) = _broadcast_unary_nz2nz_z2z_T(abs2, A, TTv) +abs{TTv}(A::SparseMatrixCSC{Complex{TTv}}) = _broadcast_unary_nz2nz_z2z_T(abs, A, TTv) +abs{TTv<:Integer}(A::SparseMatrixCSC{Complex{TTv}}) = _broadcast_unary_nz2nz_z2z_T(abs, A, Float64) +abs{TTv<:BigInt}(A::SparseMatrixCSC{Complex{TTv}}) = _broadcast_unary_nz2nz_z2z_T(abs, A, BigFloat) function conj!(A::SparseMatrixCSC) - nzvalA = A.nzval - @simd for i=1:length(nzvalA) - @inbounds nzvalA[i] = conj(nzvalA[i]) + @inbounds @simd for k in 1:nnz(A) + A.nzval[k] = conj(A.nzval[k]) end return A end -conj(A::SparseMatrixCSC) = conj!(copy(A)) - -# Operations that map nonzeros to nonzeros, and zeros to nonzeros -# Result is dense -for op in (:cos, :cosh, :acos, :sec, :csc, :cot, :acot, :sech, - :csch, :coth, :asech, :acsch, :cospi, :sinc, :cosd, - :cotd, :cscd, :secd, :acosd, :acotd, :log, :log2, :log10, - :exp, :exp2, :exp10) - @eval begin - - function ($op){Tv}(A::SparseMatrixCSC{Tv}) - B = fill($(op)(zero(Tv)), size(A)) - @inbounds for col = 1 : A.n - for j = A.colptr[col] : A.colptr[col+1]-1 - row = A.rowval[j] - nz = A.nzval[j] - B[row,col] = $(op)(nz) - end - end - return B - end - - end -end +# Operations that map both zeros and nonzeros to zeros, yielding a dense matrix +""" +Takes unary function `f` that maps both zeros and nonzeros to nonzeros, and returns a new +`Matrix{TvB}` `B` effectively generated by applying `f` to every entry in `A`. +""" +function _broadcast_unary_nz2nz_z2nz{Tv}(f::Function, A::SparseMatrixCSC{Tv}) + B = fill(f(zero(Tv)), size(A)) + @inbounds for j in 1:A.n + for k in nzrange(A, j) + i = A.rowval[k] + x = A.nzval[k] + B[i,j] = f(x) + end + end + return B +end +@_enumerate_childmethods(_broadcast_unary_nz2nz_z2nz, + log, log2, log10, exp, exp2, exp10, sinc, cospi, + cos, cosh, cosd, acos, acosd, + cot, coth, cotd, acot, acotd, + sec, sech, secd, asech, + csc, csch, cscd, acsch) ## Broadcasting kernels specialized for returning a SparseMatrixCSC