generic broadcast/broadcast! for sparse matrices, with #19372 semantics

[Sacha0]

This pull request provides (1) a new broadcast/broadcast! implementation specialized for a pair of (input) sparse matrices, and (2) a general broadcast/broadcast! implementation capable of efficiently handling an arbitrary number of sparse matrices (within compiler limits). This pull request also integrates these broadcast/broadcast! methods with the new map/map! implementation specialized for sparse matrices, and touches up a few rough edges of the latter.

The semantics of these broadcast! methods are those discussed in #19372. Notably these broadcast! methods only allocate when strictly necessary, so for example

julia> C = spzeros(5,5); A = speye(5); B = sparse(Bidiagonal(zeros(5), ones(4), true));

julia> @allocated broadcast!(*, C, A, B) # after warmup
0

Additionally, where these broadcast! methods do allocate, they do so slightly less pessimistically than the existing broadcast! methods. [1]

(Edited) For operations that yield zero when at least one argument is zero (like *), the new broadcast/broadcast! methods specialized for a pair of (input) sparse matrices achieve substantially better performance than corresponding existing broadcast/broadcast! methods (despite performing slightly more work to enable more precise allocation behavior).

(Edited) For operations that yield zero only when all arguments are zero (like +), the new broadcast/broadcast! methods specialized for a pair of (input) sparse matrices achieve performance comparable to corresponding existing broadcast/broadcast! methods. The former win in some cases and the latter in others, with runtime differences usually within ~30% in (hopefully realistic) tests.

The general broadcast/broadcast! methods achieve performance not too far behind the new methods specialized for pairs of (input) sparse matrices. The former take from 1% to 80% again the latter's runtime, with performance for small matrices falling in the middle to slower half of that range, and performance for large matrices falling in the faster quarter to faster half of that range.

Relevant to #16285 and friends as a (near-final?) step in the sparse broadcast/broadcast! overhaul. Best!

[1] broadcast/broadcast! could allocate still less pessimistically, but whether the additional code complexity would be worth the improvement is not clear. An extreme example: broadcast[!] operations that both allocate and involve one-by-one sparse matrices allocate enough storage for a dense result, though that often isn't necessary. Implementing smarter behavior for that case is straightforward but complicates the code somewhat, and broadcasting with one-by-one sparse matrices seems like a marginal case.

[stevengj]

Is there a similar test for combinations of sparse matrices and scalars, e.g. A .* α where α is a scalar?

[stevengj]

I guess this PR doesn't handle combinations of sparse matrices and scalars at all? Ideally we'd have something that worked for any such combination.

[stevengj]

My thinking was that containertype on a sparse array should return SparseArray, and promote_containertype would produce SparseArray for combinations of SparseArray and Any (scalars), but would produce Array for combinations with Array or Tuple.

[Sacha0]

I guess this PR doesn't handle combinations of sparse matrices and scalars at all? Ideally we'd have something that worked for any such combination.

Correct, and agreed. The next step is a pull request including the work in this review thread, hopefully capable of handling combinations involving more than two SparseMatrixCSCs.

My thinking was that containertype on a sparse array should return SparseArray, and promote_containertype would produce SparseArray for combinations of SparseArray and Any (scalars), but would produce Array for combinations with Array or Tuple

Yes, that's the approach in the thread mentioned above. Thanks for reviewing!

[stevengj]

probably just eltype(A.rowval)(C.m + 1), since that does a type assertion to ensure the compiler knows the type of the result, unlike convert. Though I wonder if it is better to just use typemax(eltype(A.rowval)) as the sentinel?

[Sacha0]

probably just eltype(A.rowval)(C.m + 1), since that does a type assertion to ensure the compiler knows the type of the result, unlike convert.

Does the implicit type assertion provide any benefit when type inference is happy with the explicit converts? Explicit converts seem easier to understand. I suppose rowsentinelA::eltype(A.rowval) = C.m + 1 might provide the best of both worlds?

Though I wonder if it is better to just use typemax(Int) as the sentinal in all cases?

Using typemax(Int) as the sentinel might cause an issue if a rowval vector's element type is narrower than Int?

Thanks!

[Sacha0]

I suppose rowsentinelA::eltype(A.rowval) = C.m + 1 might provide the best of both worlds?

A further thought: A few assertions like Ak::eltype(A.rowval) might allow us to simplify some expressions, e.g. Ak += one(Ak) -> Ak += 1. So perhaps adopting that style uniformly in these methods would be advantageous?

[stevengj]

It seems like you might need more tests to get coverage of all your branches, e.g. the cases where A or B has only one row and/or column?

[Sacha0]

It seems like you might need more tests to get coverage of all your branches, e.g. the cases where A or B has only one row and/or column?

Do you mean in the result tests, or in the allocation tests?

If the former, existing tests should thoroughly cover the code specialized for two SparseMatrixCSCs, and if I'm not mistaken (though quite possibly I am) the new tests should cover all branches of the general code.

If the latter, yes, for allocation I only cover the no-singleton-dims and second-dim-singleton code. Would you like me to expand the allocation tests to also include the first-dim-singleton code?

Thanks!

[stevengj]

I meant in the results tests; I wasn't sure if we had broadcast tests for all of those cases.

[Sacha0]

I meant in the results tests; I wasn't sure if we had broadcast tests for all of those cases.

I believe we do, distributed piecemeal over test/sparse/sparse.jl presently. (Consolidating those tests would be worthwhile at some point.) Thanks again!

[martinholters]

Hm, before:

julia> eltype(sparse(zeros(Real, 0, 0)) + sparse(zeros(Real, 0, 0)))
Real

After:

julia> eltype(sparse(zeros(Real, 0, 0)) + sparse(zeros(Real, 0, 0)))
Any

I found the old behavior rather handy, but there is no way of knowing that +(::Real, ::Real) will always return a Real, save of the old, dreaded, unscalable, hand-written promote_ops, right? So I probably have to adjust to the new behavior, or am I missing something we could do here?

[nalimilan]

In theory that should work if all + methods on reals return a real. For example, in this very simple case:

julia> f(x, y) = x
f (generic function with 1 method)

julia> Core.Inference.return_type(f, Tuple{Real, Real})
Real

But for + there must be at least one method which is not correctly inferred:

julia> last(Base.return_types(+, Tuple{Real, Real}))
Any

I guess we should find it and fix it. But unfortunately any package adding a non-inferrable function will break this again.

[martinholters]

I guess we should find it and fix it.

The offenders are +{T<:AbstractFloat}(x::Bool, y::T)::promote_type(Bool,T) and +(x::Number, y::Number), so the promotion mechanism seems to be the culprit.

But unfortunately any package adding a non-inferrable function will break this again.

... or adding a WeirdReal <: Real, where +(::WeirdReal, ::WeirdReal) does not give a Real. The prospect of this definition breaking my code although I never used WeirdReal suffices to make me not rely on the element type of empty matrices being correct.

[martinholters]

Ah, but for sparse matrices, the problem is a bit more grave, as it also hits the empty as in "no stored elements" case, not only the empty as in "at least one dimension is zero" case:

julia> spzeros(Real,1,1) + spzeros(Real,1,1)
1×1 sparse matrix with 0 Any nonzero entries

The element type should probably be typeof(zero(Real)+zero(Real)) (==Int64), that would also be in line with the dense matrices, e.g. compare:

julia> eltype(broadcast(+, zeros(Real,1,1), zeros(Real,1,1)))
Int64

julia> eltype(broadcast(+, sparse(zeros(Real,1,1)), sparse(zeros(Real,1,1))))
Any

@Sacha0, do you think we can update the sparse matrix code to behave like the dense array code?

[nalimilan]

But:

julia> zeros(Real,1,1) + zeros(Real,1,1)
1×1 Array{Real,2}:
 0

For broadcast, promote_op used to preserve the input element type if it was a supertype of the output element type. Was this removed intentionally?

Also, promote_type(Bool, T) seems to always return Real when passed a Real according to return_type, so it doesn't seem to be the problem. And +(x::Number, y::Number) throws an error (i.e. a Union{} return type), which shouldn't be an issue. Why do you think they are the offenders?

[martinholters]

IMHO, the non-broadcast + and - on array are a bit of a hack concerning their return-type. I'm not sure we want to duplicate that, but it would indeed fix the immediate issue.

For broadcast, promote_op used to preserve the input element type if it was a supertype of the output element type. Was this removed intentionally?

Probably, otherwise broadcast(x -> round(Int, x), Real[1, 1.5]) would return a Vector{Real}, which would be unfortunate, too.

promote_type(Bool, T) seems to always return Real when passed a Real according to return_type

But that does not necessarily imply that for any T<:Real, promote_type(Bool, T)<:Real, which inference would need to figure out.

And +(x::Number, y::Number) throws an error

Errr, +{T<:Number}(x::T, y::T) does error, +(x::Number, y::Number) = +(promote(x,y)...) should not error.

[nalimilan]

Probably, otherwise broadcast(x -> round(Int, x), Real[1, 1.5]) would return a Vector{Real}, which would be unfortunate, too.

Sure, but both choices have their drawbacks, since the stricter type means some Real won't fit into the array, which you could have expected. Anyway, we need to make this decision a conscious and consistent choice everywhere.

But that does not necessarily imply that for any T<:Real, promote_type(Bool, T)<:Real, which inference would need to figure out.

What do you mean? At least this gives the expected result:

julia> f{T}(::Type{T}) = promote_type(Bool, T)
f (generic function with 1 method)

julia> Core.Inference.return_type(f, Tuple{Type{Real}})
Type{Real}

Errr, +{T<:Number}(x::T, y::T) does error, +(x::Number, y::Number) = +(promote(x,y)...) should not error.

Indeed, that's the problematic definition. It seems that even when no promotion rule exists, inference is not able to figure out that the function cannot return:

julia> f{S,T}(x::S, y::T) = (U=promote_type(S,T); convert(U, x) + convert(U, y))
f (generic function with 1 method)

julia> type Num <: Number
       end

julia> @code_warntype f(Num(), 1)
Variables:
  #self#::#f
  x::Num
  y::Int64
  U::Type{T}

Body:
  begin 
      $(Expr(:inbounds, false))
      # meta: location promotion.jl promote_type 140
      # meta: location promotion.jl promote_result 178
      # meta: location promotion.jl promote_to_supertype 187
      SSAValue(1) = Num
      SSAValue(2) = Int64
      # meta: pop location
      # meta: pop location
      # meta: pop location
      $(Expr(:inbounds, :pop))
      U::Type{T} = (Base.throw)((Base.ErrorException)(((Core.getfield)((Core.getfield)(Base.Main,:Base)::Any,:string)::Any)("no promotion exists for ",SSAValue(1)," and ",SSAValue(2))::Any)::ErrorException)::Union{}
      return ((Main.convert)(U::Type{T},x::Num)::Any + (Main.convert)(U::Type{T},y::Int64)::Any)::Any
  end::Any

Here it seems that U should be inferred as Union{}.

[Sacha0]

@Sacha0, do you think we can update the sparse matrix code to behave like the dense array code?

Yes, my intention was to change the sparse broadcast promotion mechanism in this pull request to that provided in #19421 prior to merging. #19576 (thanks!) should have fixed this discrepancy? (Edit: i.e. the _promote_op versus _default_eltype-... discrepancies, not the others.) Best!

[Sacha0]

Hm, before:

julia> eltype(sparse(zeros(Real, 0, 0)) + sparse(zeros(Real, 0, 0)))
Real
After:

julia> eltype(sparse(zeros(Real, 0, 0)) + sparse(zeros(Real, 0, 0)))
Any
I found the old behavior rather handy, but there is no way of knowing that +(::Real, ::Real) will always return a Real, save of the old, dreaded, unscalable, hand-written promote_ops, right? So I probably have to adjust to the new behavior, or am I missing something we could do here?

The former methods only supported broadcasting over one or two input sparse matrices, and determined output element type via the unary and binary promote_ops (through promote_eltype_op). The new methods support broadcasting over any number of input sparse matrices, and so instead determine output element type via the more general broadcast promotion mechanism (formerly the _default_eltype-first-generator-zip mechanism, and now with #19421 and #19576 the _promote_op(f, typestuple(As...)) mechanism).

We could continue to use promote_op when broadcasting over one or two input sparse matrices at the cost of inconsistency with cases involving more than two input sparse matrices. (Is there a fundamental reason the promote_ops cannot be extended programmatically to more than two type arguments?) Best!

[Sacha0]

What do you mean? At least this gives the expected result:

julia> f{T}(::Type{T}) = promote_type(Bool, T)
f (generic function with 1 method)

julia> Core.Inference.return_type(f, Tuple{Real,Real})
Union{}

@nalimilan, could you clarify this example? Thanks!

[nalimilan]

@Sacha0 Good catch, I copied/pasted an incorrect output. This should have been Core.Inference.return_type(f, Tuple{Type{Real}}) -> Type{Real} (I've edited my post).

We could continue to use promote_op when broadcasting over one or two input sparse matrices at the cost of inconsistency with cases involving more than two input sparse matrices. (Is there a fundamental reason the promote_ops cannot be extended programmatically to more than two type arguments?) Best!

We should definitely get rid of any inconsistencies. I don't think there's a strong reason which prevents extending promote_op to any number of arguments. Though I think the plan is rather to get rid of promote_op and always rely on inference instead. In most cases this should give identical results, except for inference failures like the case at hand, and for the question of whether to preserve abstract types or not.

[martinholters]

I think, in a nutshell, promote_op extends _promote_op with some special provisions suitable for mathematical operations. Now, on the one hand, broadcast does not know what kind of operation it is broadcasting, so using promote_op seems a bit arbitrary. On the other hand, arguably sparse matrices are primarily used for mathematical operations, which might justify using promote_op here (but would lead to inconsistencies with Array). I'm torn whether this is a good idea.

Please also see #19595 which I have opened to continue the general discussion.