Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

more generic broadcast over a single sparse matrix #19239

Merged
merged 1 commit into from
Nov 21, 2016
Merged

more generic broadcast over a single sparse matrix #19239

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
134 changes: 52 additions & 82 deletions base/sparse/sparsematrix.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1398,31 +1398,19 @@ end

sparse(S::UniformScaling, m::Integer, n::Integer=m) = speye_scaled(S.λ, m, n)

## Unary arithmetic and boolean operators

"""
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 `broadcast(::typeof(fc), A::SparseMatrixCSC) = fp(fc, A)`.
"""
macro _enumerate_childmethods(fp, fcs...)
fcexps = Expr(:block)
for fc in fcs
push!(fcexps.args, :( $(esc(:broadcast))(::typeof($(esc(fc))), A::SparseMatrixCSC) = $(esc(fp))($(esc(fc)), A) ) )
end
return fcexps
end
## Broadcast operations involving a single sparse matrix and possibly broadcast scalars

# 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))
function broadcast{Tf}(f::Tf, A::SparseMatrixCSC)
fofzero = f(zero(eltype(A)))
fpreszero = fofzero == zero(fofzero)
return fpreszero ? _broadcast_zeropres(f, A) : _broadcast_notzeropres(f, fofzero, A)
end
"Returns a `SparseMatrixCSC` storing only the nonzero entries of `broadcast(f, Matrix(A))`."
function _broadcast_zeropres{Tf}(f::Tf, A::SparseMatrixCSC)
Bcolptr = similar(A.colptr, A.n + 1)
Browval = similar(A.rowval, nnz(A))
Bnzval = similar(A.nzval, Base.Broadcast.promote_eltype_op(f, A), nnz(A))
Bk = 1
@inbounds for j in 1:A.n
Bcolptr[j] = Bk
Expand All @@ -1440,69 +1428,51 @@ function _broadcast_unary_nz2z_z2z_T{TvA,TiA,TvB}(f::Function, A::SparseMatrixCS
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)
broadcast(::typeof(real), A::SparseMatrixCSC) = copy(A)
broadcast{Tv,Ti}(::typeof(imag), A::SparseMatrixCSC{Tv,Ti}) = spzeros(Tv, Ti, A.m, A.n)
broadcast{TTv}(::typeof(real), A::SparseMatrixCSC{Complex{TTv}}) = _broadcast_unary_nz2z_z2z_T(real, A, TTv)
broadcast{TTv}(::typeof(imag), 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{TvA,TiA,Tf<:Function}(f::Tf, A::SparseMatrixCSC{TvA,TiA})
Bcolptr = Vector{TiA}(A.n + 1)
Browval = Vector{TiA}(nnz(A))
copy!(Bcolptr, 1, A.colptr, 1, A.n + 1)
copy!(Browval, 1, A.rowval, 1, nnz(A))
Bnzval = broadcast(f, A.nzval)
resize!(Bnzval, nnz(A))
return SparseMatrixCSC(A.m, A.n, Bcolptr, Browval, Bnzval)
end
@_enumerate_childmethods(_broadcast_unary_nz2nz_z2z,
log1p, expm1, abs, abs2, conj)
function conj!(A::SparseMatrixCSC)
@inbounds @simd for k in 1:nnz(A)
A.nzval[k] = conj(A.nzval[k])
end
return A
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, constructs a new
`Matrix{TvB}` `B`, populates all entries of `B` with the result of a single `f(one(zero(Tv)))`
call, and then, for each stored entry in `A`, calls `f` on the entry's value and stores the
result in the corresponding location in `B`.
Returns a (dense) `SparseMatrixCSC` with `fillvalue` stored in place of each unstored
entry in `A` and `f(A[i,j])` stored in place of each stored entry `A[i,j]` 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
function _broadcast_notzeropres{Tf}(f::Tf, fillvalue, A::SparseMatrixCSC)
nnzB = A.m * A.n
# Build structure
Bcolptr = similar(A.colptr, A.n + 1)
copy!(Bcolptr, 1:A.m:(nnzB + 1))
Browval = similar(A.rowval, nnzB)
for k in 1:A.m:(nnzB - A.m + 1)
copy!(Browval, k, 1:A.m)
end
# Populate values
Bnzval = fill(fillvalue, nnzB)
@inbounds for (j, jo) in zip(1:A.n, 0:A.m:(nnzB - 1)), k in nzrange(A, j)
Bnzval[jo + A.rowval[k]] = f(A.nzval[k])
end
return B
# NOTE: Combining the fill call into the loop above to avoid multiple sweeps over /
# nonsequential access of Bnzval does not appear to improve performance
return SparseMatrixCSC(A.m, A.n, Bcolptr, Browval, Bnzval)
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)

# Cover common broadcast operations involving a single sparse matrix and one or more
# broadcast scalars.
#
# TODO: The minimal snippet below is not satisfying: A better solution would achieve
# the same for (1) all broadcast scalar types (Base.Broadcast.containertype(x) == Any?) and
# (2) any combination (number, order, type mixture) of broadcast scalars.
#
broadcast{Tf}(f::Tf, x::Union{Number,Bool}, A::SparseMatrixCSC) = broadcast(y -> f(x, y), A)
broadcast{Tf}(f::Tf, A::SparseMatrixCSC, y::Union{Number,Bool}) = broadcast(x -> f(x, y), A)
# NOTE: The following two method definitions work around #19096. These definitions should
# be folded into the two preceding definitions on resolution of #19096.
broadcast{Tf,T}(f::Tf, ::Type{T}, A::SparseMatrixCSC) = broadcast(y -> f(T, y), A)
broadcast{Tf,T}(f::Tf, A::SparseMatrixCSC, ::Type{T}) = broadcast(x -> f(x, T), A)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The four method definitions above for handling common combinations of sparse matrices with broadcast scalars is not satisfying: A better solution would achieve the same for (1) all broadcast scalar types (Base.Broadcast.containertype(x) == Any?) and (2) any combination (number, order, type mixture) of broadcast scalars. One might achieve (2) via a generated function. But how one might simultaneously and gracefully achieve (1) eludes me. Thoughts much appreciated. Thanks!

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think you need a generated function. Maybe we can brainstorm with @pabloferz about this at some point soon? Maybe start a discourse thread? @pabloferz, are you on discourse.julialang.org yet?

In any case, it seems like that should be separate from this PR, which is only about the single-argument case.

Copy link
Contributor

@pabloferz pabloferz Nov 12, 2016 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Al least for this case you might want to add another trait layer. E.g.

import Base.Broadcast.containertype
broadcast(f, A, B::SparseMatrixCSC) = broadcast_sp(f, containertype(A), containertype(B), A, B)
broadcast(f, A::SparseMatrixCSC, B) = broadcast_sp(f, containertype(A), containertype(B), A, B)
broadcast_sp(f, ::Type{Any}, ::Type{Array} A, B) = broadcast(x -> f(A, x), B)
broadcast_sp(f, ::Type{Array}, ::Type{Any} A, B) = broadcast(x -> f(x, B), A)
broadcast_sp(f, ::Type, ::Type, A, B) = sparse(Base.Broadcast.broadcast_c(f, Array, A, B))

@stevengj I'm already on discourse, so feel free to ping me (@pabloferz) if you'd like to start a thread overe there.

Copy link
Member Author

@Sacha0 Sacha0 Nov 14, 2016 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The following snippet works for arbitrary combinations of broadcast scalars with a single sparse matrix, but the (y...) = f(a, y...) kills inference.

import Base.Broadcast: broadcast_c, containertype, promote_containertype

containertype{T<:SparseMatrixCSC}(::Type{T}) = SparseMatrixCSC
# Combinations of sparse arrays with broadcast scalars should yield sparse arrays
promote_containertype(::Type{Any}, ::Type{SparseMatrixCSC}) = SparseMatrixCSC
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Any}) = SparseMatrixCSC
# Combinations of sparse arrays with anything else should fall back to dense array methods
promote_containertype(::Type{Array}, ::Type{SparseMatrixCSC}) = Array
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Array}) = Array
promote_containertype(::Type{Tuple}, ::Type{SparseMatrixCSC}) = Array
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Tuple}) = Array

broadcast_c(f, ::Type{SparseMatrixCSC}, A) = broadcast(f, A)
broadcast_c(f, ::Type{SparseMatrixCSC}, A, Bs...) = broadcast_sp(f, containertype(A), containertype(Bs...), A, Bs...)
broadcast_sp(f, ::Type{Any}, ::Type{SparseMatrixCSC}, a, Bs...) = broadcast_c((y...) -> f(a, y...), SparseMatrixCSC, Bs...)
broadcast_sp(f, ::Type{SparseMatrixCSC}, ::Type{Any}, A, bs...) = broadcast_c(x -> f(x, bs...), SparseMatrixCSC, A)
broadcast_sp(f, ::Type{SparseMatrixCSC}, ::Type{SparseMatrixCSC}, A, Bs...) = broadcast_c(f, Array, A, Bs...)

Copy link
Contributor

@pabloferz pabloferz Nov 14, 2016 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A possible solution would be to have (additionally to what you have above)

EDIT: There was a typo. This works with one SparseMatrixCSC and up to ten scalars but won't be inferable with two or more SparseMatrixCSCs.

EDIT 2: Added a method, and it appears to be inferable with up to eight arguments (I'd say that should be enough).

immutable FArgs{F,T}
    f::F
    args::T
end
(f::FArgs)(x...) = f.f(f.args..., x...)

broadcast_sp(f, ::Type{Any}, ::Type{SparseMatrixCSC}, a, Bs...) = broadcast_c(FArgs(f,(a,)), SparseMatrixCSC, Bs...)
broadcast_sp(f::FArgs, ::Type{Any}, ::Type{SparseMatrixCSC}, a, Bs...) = broadcast_c(FArgs(f.f,(f.args...,a)), SparseMatrixCSC, Bs...)
broadcast_sp(f::FArgs, ::Type{SparseMatrixCSC}, ::Type{Any}, A, bs...) = broadcast_c(x -> f.f(f.args..., x, bs...), SparseMatrixCSC, A)
broadcast_sp(f::FArgs, ::Type{SparseMatrixCSC}, ::Type{SparseMatrixCSC}, A, Bs...) = broadcast_c(f.f, Array, f.args..., A, Bs...)

Copy link
Member Author

@Sacha0 Sacha0 Nov 14, 2016 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The hacky snippet below takes a different tack. In principle it should handle any combination of broadcast scalars and SparseMatrixCSCs; in practice type inference fails when (1) the number of arguments exceeds foursix or (2) more than one SparseMatrixCSC is involved. (edited:)

import Base: tail, front
import Base.Broadcast: broadcast_c, containertype, promote_containertype, broadcast_indices

broadcast_indices(::Type{SparseMatrixCSC}, A) = indices(A)
containertype{T<:SparseMatrixCSC}(::Type{T}) = SparseMatrixCSC
# Combinations of sparse arrays with broadcast scalars should yield sparse arrays
promote_containertype(::Type{Any}, ::Type{SparseMatrixCSC}) = SparseMatrixCSC
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Any}) = SparseMatrixCSC
# Combinations of sparse arrays with anything else should fall back to dense array methods
promote_containertype(::Type{Array}, ::Type{SparseMatrixCSC}) = Array
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Array}) = Array
promote_containertype(::Type{Tuple}, ::Type{SparseMatrixCSC}) = Array
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Tuple}) = Array

broadcast_c(f, ::Type{SparseMatrixCSC}, As...) = foo((x, y) -> f(y...), (), As...)

foo(f, arrays, nextarg, remargs...) = foo((x, y) -> f(x, (nextarg, y...)), arrays, remargs...)
foo(f, arrays, nextarg::SparseMatrixCSC, remargs...) =
    foo((x, y) -> f(front(x), (last(x), y...)), (arrays..., nextarg), remargs...)

foo(f, arrays, nextarg) = bar(x -> f(x, (nextarg,)), arrays...)
foo(f, arrays, nextarg::SparseMatrixCSC) = bar(x -> f(front(x), (last(x),)), arrays..., nextarg)

bar(f, A) = broadcast(x -> f((x,)), A)
bar(f, A, B) = broadcast_c((x,y) -> f((x,y)), Array, A, B)
bar(f, A, B, Cs...) = broadcast_c((x,y,zs...) -> f((x,y,zs...)), Array, A, B, Cs...)

I'm sure a better approach exists, possibly e.g. inverting the call chain. Having a look at @pabloferz's suggestion now. Best!

Copy link
Member Author

@Sacha0 Sacha0 Nov 15, 2016 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Below are two evolutions of @pabloferz's suggestion. The first variation retains the helper type, which distills to a stack for leading scalar arguments. The second variation nixes the helper type and instead uses a tuple for the stack.

Both variations reduce the mutual recursion between broadcast_c and broadcast_sp to recursion simply in broadcast_sp, and punt on argument lists involving more than one SparseMatrixCSC.

Manually extending this approach to handle argument lists involving two (or some larger small integer N) SparseMatrixCSCs seems straightforward, but incurs rapid code growth with N. How to extend this approach to an arbitrary number of SparseMatrixCSCs isn't clear to me (see followup post). But perhaps N = 2 or N = 3 would be good enough.

I will play with the N = 2 case of this approach next, and then attempt to improve (inference of) the more general approach I posted above.

Code common to both variations

import Base: tail
import Base.Broadcast: broadcast_c, containertype, promote_containertype, broadcast_indices

broadcast_indices(::Type{SparseMatrixCSC}, A) = indices(A)
containertype{T<:SparseMatrixCSC}(::Type{T}) = SparseMatrixCSC
# Combinations of sparse arrays with broadcast scalars should yield sparse arrays
promote_containertype(::Type{Any}, ::Type{SparseMatrixCSC}) = SparseMatrixCSC
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Any}) = SparseMatrixCSC
# Combinations of sparse arrays with anything else should fall back to dense array methods
promote_containertype(::Type{Array}, ::Type{SparseMatrixCSC}) = Array
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Array}) = Array
promote_containertype(::Type{Tuple}, ::Type{SparseMatrixCSC}) = Array
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Tuple}) = Array

Variation retaining helper type

immutable FuncArgs{F,T}; func::F; args::T; end

# Entrypoint to broadcast_sp
broadcast_c(f, ::Type{SparseMatrixCSC}, A, Bs...) =
    broadcast_sp(FuncArgs(f, ()), containertype(A), containertype(Bs...), A, Bs...)
# Endpoints for argument lists containing a single SparseMatrixCSC
broadcast_sp(fa, ::Type{Any}, ::Type{SparseMatrixCSC}, a, B) = broadcast(x -> fa.func(fa.args..., a, x), B)
broadcast_sp(fa, ::Type{SparseMatrixCSC}, ::Type{Any}, A, bs...) = broadcast(x -> fa.func(fa.args..., x, bs...), A)
# Endpoint for argument lists containing more than one SparseMatrixCSC (punt)
broadcast_sp(fa, ::Type{SparseMatrixCSC}, ::Type{SparseMatrixCSC}, A, Bs...) = broadcast_c(fa.func, Array, fa.args..., A, Bs...)
# Case requiring further reduction
broadcast_sp(fa, ::Type{Any}, ::Type{SparseMatrixCSC}, a, B, Cs...) =
    broadcast_sp(FuncArgs(fa.func, (fa.args..., a)), containertype(B), containertype(Cs...), B, Cs...)

Variation nixing helper type

# Entrypoint to broadcast_sp
broadcast_c(f, ::Type{SparseMatrixCSC}, A, Bs...) =
    broadcast_sp(f, (), containertype(A), containertype(Bs...), A, Bs...)
# Endpoints for argument lists containing a single SparseMatrixCSC
broadcast_sp(f, as, ::Type{Any}, ::Type{SparseMatrixCSC}, b, C) = broadcast(x -> f(as..., b, x), C)
broadcast_sp(f, as, ::Type{SparseMatrixCSC}, ::Type{Any}, B, cs...) = broadcast(x -> f(as..., x, cs...), B)
# Endpoint for argument lists containing more than one SparseMatrixCSC (punt)
broadcast_sp(f, as, ::Type{SparseMatrixCSC}, ::Type{SparseMatrixCSC}, B, Cs...) = broadcast_c(f, Array, as..., B, Cs...)
# Case requiring further reduction
broadcast_sp(f, as, ::Type{Any}, ::Type{SparseMatrixCSC}, b, C, Ds...) =
    broadcast_sp(f, (as..., b), containertype(C), containertype(Ds...), C, Ds...)

Best!

Copy link
Member Author

@Sacha0 Sacha0 Nov 17, 2016 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

New versions capable (in principle) of handling any number of SparseMatrixCSCs follow below. None circumvent #19096. TL;DR The first version (A) is the best so far, perhaps good enough to incorporate for now? Thoughts? Any and all pointers much appreciated, particularly towards better understanding of why type inference works/fails where it does.

Code common to all versions:

import Base: front, tail
import Base.Broadcast: broadcast_c, containertype, promote_containertype, broadcast_indices

# Broadcast container type promotion for SparseMatrixCSCs
containertype{T<:SparseMatrixCSC}(::Type{T}) = SparseMatrixCSC
broadcast_indices(::Type{SparseMatrixCSC}, A) = indices(A)
# Combinations of sparse arrays with broadcast scalars should yield sparse arrays
promote_containertype(::Type{Any}, ::Type{SparseMatrixCSC}) = SparseMatrixCSC
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Any}) = SparseMatrixCSC
# Combinations of sparse arrays with anything else should fall back to dense array
promote_containertype(::Type{Array}, ::Type{SparseMatrixCSC}) = Array
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Array}) = Array
promote_containertype(::Type{Tuple}, ::Type{SparseMatrixCSC}) = Array
promote_containertype(::Type{SparseMatrixCSC}, ::Type{Tuple}) = Array

(A) For the following version, type inference remains happy when the argument list (1) contains a single sparse matrix and up to six broadcast scalars, (2) contains two sparse matrices and up to twelve broadcast scalars (more depending on position, strangely), or (3) up to six sparse matrices as long as no broadcast scalars appear. This behavior is a mystery to me, insight much appreciated. This version seems the best in practice at the moment.

# closeoverscalars steps through the argument list. For each scalar argument,
# closeoverscalars generates a closure incorporating that scalar into the present version
# of the broadcast function, removes that scalar from the argument list, and steps to the
# next argument. For each matrix argument, closeoverscalars generates a closure which
# will later funnel that matrix argument to the right slot in the original broadcast
# function, collects that matrix in a tuple of matrices over which to eventually broadcast,
# removes that matrix from the argument list, and steps to the next argument. After
# traversing the argument list, closeoverscalars returns a closure incorporating
# all broadcast scalars in the argument list followed by the argument list's sparse matrices.
#
# Entrypoint to closeoverscalars
broadcast_c(f, ::Type{SparseMatrixCSC}, args...) = broadcast(closeoverscalars(f, args)...)
closeoverscalars(f, args) = closeoverscalars((x, y) -> f(y...), (), args...)
# Recursion cases for closeoverscalars
closescalar(f, s) = (x, y) -> f(x, (s, y...))
passmatrix(f) = (x, y) -> f(front(x), (last(x), y...))
closeoverscalars(f, mats, nextarg, remargs...) = closeoverscalars(closescalar(f, nextarg), mats, remargs...)
closeoverscalars(f, mats, nextarg::SparseMatrixCSC, remargs...) = closeoverscalars(passmatrix(f), (mats..., nextarg), remargs...)
# Base cases for closescalars
closelastscalar(f, s) = (x...) -> f(x, (s,))
passlastmatrix(f) = (x...) -> f(front(x), (last(x),))
closeoverscalars(f, mats, lastarg) = (closelastscalar(f, lastarg), mats...)
closeoverscalars(f, mats, lastarg::SparseMatrixCSC) = (passlastmatrix(f), mats..., lastarg)

Edit: Partial evaluation not correct. (B) For the following version (a descendant of @pabloferz's suggestion), type inference remains happy when the argument list (1) contains up to nine members including either one or two sparse matrices, or (2) contains three sparse matrices and no broadcast scalars. This behavior is also a mystery to me, insight much appreciated.

# broadcast_sp steps through the argument list, collecting each sequence of broadcast scalars
# between one sparse matrix (or start or end of the argument list) and the next in a tuple,
# storing the sequence of such tuples in stups ("scalar-tuples"), and storing the sparse
# matrices in mats. After traversing the argument list, broadcast_sp calls closeoverscalars
# which, given stups, mats, and the broadcast function, generates a closure incorporating
# all broadcast scalars in the argument list. broadcast_sp then passes that closure along
# with the argument list's sparse matrices on to broadcast.
#
# Entrypoint to broadcast_sp
broadcast_c(fn, ::Type{SparseMatrixCSC}, A, Bs...) =
    broadcast_sp(fn, ((),), (), containertype(A), containertype(Bs...), A, Bs...)
# Recursion cases for broadcast_sp
broadcast_sp(fn, stups, mats, ::Type{Any}, ::Type{SparseMatrixCSC}, a, B, Cs...) =
    broadcast_sp(fn, (front(stups)..., (last(stups)..., a)), mats, containertype(B), containertype(Cs...), B, Cs...)
broadcast_sp(fn, stups, mats, ::Type{SparseMatrixCSC}, ::Type{SparseMatrixCSC}, A, B, Cs...) =
    broadcast_sp(fn, (stups..., ()), (mats..., A), containertype(B), containertype(Cs...), B, Cs...)
# Base cases for broadcast_sp
broadcast_sp(fn, stups, mats, ::Type{SparseMatrixCSC}, ::Type{SparseMatrixCSC}, A, B) =
    broadcast(closeoverscalars((x...) -> fn(front(front(x))..., last(stups)..., last(front(x)), last(x)), front(stups)), mats..., A, B)
broadcast_sp(fn, stups, mats, ::Type{SparseMatrixCSC}, ::Type{Any}, A, bs...) =
    broadcast(closeoverscalars((x...) -> fn(front(x)..., last(stups)..., last(x), bs...), front(stups)), mats..., A)
broadcast_sp(fn, stups, mats, ::Type{Any}, ::Type{SparseMatrixCSC}, a, B) =
    broadcast(closeoverscalars((x...) -> fn(front(x)..., last(stups)..., a, last(x)), front(stups)), mats..., B)
# Recursively close over scalars
closeoverscalars(fn, stups) = closeoverscalars((x...) -> fn(front(x)..., last(stups)..., last(x)), front(stups))
closeoverscalars(fn, ::Tuple{}) = fn

Edit: Partial evaluation not correct. (C) The following version is a simpler realization of the preceding concept, but type inference does not seem to like it (nor a similar version using varags for groupargs).

# Entrypoint
broadcast_c(f, ::Type{SparseMatrixCSC}, args...) =
    (ga = groupargs(args); broadcast(closeoverscalars(f, ga[1]), ga[2]...))
# Entrypoint for groupargs
groupargs(args) = groupargs(((),), (), first(args), tail(args))
# Recursion cases for groupargs
groupargs(stups, mats, nextarg, remargs) =
    groupargs((front(stups)..., (last(stups)..., nextarg)), mats, first(remargs), tail(remargs))
groupargs(stups, mats, nextarg::SparseMatrixCSC, remargs) =
    groupargs((stups..., ()), (mats..., nextarg), first(remargs), tail(remargs))
# Base cases for groupargs
groupargs(stups, mats, lastarg, ::Tuple{}) = ((front(stups)..., (last(stups)..., lastarg)), mats)
groupargs(stups, mats, lastarg::SparseMatrixCSC, ::Tuple{}) = ((stups..., ()), (mats..., lastarg))
# Entrypoint and recursion and base cases for closeoverscalars
closeoverscalars(f, stups) = _closeoverscalars((x...) -> f(x..., last(stups)...), front(stups))
_closeoverscalars(f, stups) = _closeoverscalars((x...) -> f(front(x)..., last(stups)..., last(x)), front(stups))
_closeoverscalars(f, ::Tuple{}) = f

Best!

Copy link
Contributor

@pabloferz pabloferz Nov 18, 2016 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Sacha0 Inspired on your approach above that evolved from my suggestion I was able to write it in such a way that it is inferable up to eight arguments irrespectively of the arguments types. The maximum number of arguments, so that it keeps inferable, I believe, is a limitation of the inference mechanism, but I haven't looked this into detail.

The approach would be your 'common code' above plus

# Entrypoint to broadcast_sp
broadcast_c(f, ::Type{SparseMatrixCSC}, A, Bs...) =
    broadcast_sp(f, (), nothing, (), containertype(A), containertype(Bs...), A, Bs...)
# Recursion cases for broadcast_sp
broadcast_sp(f, ftup, ::Void, ttup, ::Type{Any}, ::Type{SparseMatrixCSC}, a, B, Cs...) =
    broadcast_sp(f, (ftup..., a), nothing, ttup, containertype(B), containertype(Cs...), B, Cs...)
broadcast_sp(f, ftup, ::Void, ttup, ::Type{SparseMatrixCSC}, ::Type{SparseMatrixCSC}, A, B, Cs...) =
    broadcast_sp(f, ftup, A, ttup, containertype(B), containertype(Cs...), B, Cs...)
broadcast_sp(f, ftup, A, ttup, ::Type{Any}, ::Type{SparseMatrixCSC}, a, B, Cs...) =
    broadcast_sp(f, ftup, A, (ttup..., a), containertype(B), containertype(Cs...), B, Cs...)
# Base cases for broadcast_sp
broadcast_sp(f, ftup, ::Void, ttup, ::Type{Any}, ::Type{SparseMatrixCSC}, a, B) =
    broadcast(x -> f(ftup..., a, x), B)
broadcast_sp(f, ftup, ::Void, ttup, ::Type{SparseMatrixCSC}, ::Type{Any}, A, bs...) =
    broadcast(x -> f(ftup..., x, bs...), A)
broadcast_sp(f, ftup, ::Void, ttup, ::Type{SparseMatrixCSC}, ::Type{SparseMatrixCSC}, A, B) =
    broadcast((x, y) -> f(ftup..., x, y), A, B)
broadcast_sp(f, ftup, A, ttup, ::Type{Any}, ::Type{SparseMatrixCSC}, a, B) =
    broadcast((x, y) -> f(ftup..., x, ttup..., a, y), A, B)
broadcast_sp(f, ftup, A, ttup, ::Type{SparseMatrixCSC}, ::Type{Any}, B, cs...) =
    broadcast((x, y) -> f(ftup..., x, ttup..., y, cs...), A, B)
broadcast_sp(f, ftup, A, ttup, ::Type{SparseMatrixCSC}, ::Type{SparseMatrixCSC}, B, Cs...) =
    broadcast_c(f, Array, ftup..., A, ttup..., B, Cs...)

NOTE: When using the dot syntax, given that the parser can already reduce the effective number of arguments passed to broadcast the number of arguments can be even more.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussion continued on discourse. Best!


# TODO: The following definitions should be deprecated.
ceil{To}(::Type{To}, A::SparseMatrixCSC) = ceil.(To, A)
floor{To}(::Type{To}, A::SparseMatrixCSC) = floor.(To, A)
trunc{To}(::Type{To}, A::SparseMatrixCSC) = trunc.(To, A)
round{To}(::Type{To}, A::SparseMatrixCSC) = round.(To, A)

# TODO: More appropriate location?
conj!(A::SparseMatrixCSC) = (broadcast!(conj, A.nzval, A.nzval); A)


## Broadcasting kernels specialized for returning a SparseMatrixCSC
Expand Down
12 changes: 4 additions & 8 deletions test/sparse/sparse.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1627,14 +1627,10 @@ let A = sparse(UInt32[1,2,3], UInt32[1,2,3], [1.0,2.0,3.0])
@test A[1,1:3] == A[1,:] == [1,0,0]
end

# Check that `broadcast` methods specialized for unary operations over
# `SparseMatrixCSC`s are called. (Issue #18705.)
let
A = spdiagm(1.0:5.0)
@test isa(sin.(A), SparseMatrixCSC) # representative for _unary_nz2z_z2z class
@test isa(abs.(A), SparseMatrixCSC) # representative for _unary_nz2nz_z2z class
@test isa(exp.(A), Array) # representative for _unary_nz2nz_z2nz class
end
# Check that `broadcast` methods specialized for unary operations over `SparseMatrixCSC`s
# are called. (Issue #18705.) EDIT: #19239 unified broadcast over a single sparse matrix,
# eliminating the former operation classes.
@test isa(sin.(spdiagm(1.0:5.0)), SparseMatrixCSC)

# 19225
let X = sparse([1 -1; -1 1])
Expand Down