Extending in-place broadcast (broadcast!) for user types

[tkoolen]

#23939 introduced a new broadcast/broadcast! API and documentation. Although I think the changes are great for broadcast, I think the story for broadcast! is a little weaker.

First, I'd like to request a bit more documentation on extending broadcast! for user types. In what follows I'm assuming that the idea is to simply add methods to broadcast!. Even if that's all, a sentence stating that that is the case would be appreciated.

Second (and maybe this can be handled by more documentation), how are users meant to extend broadcast! when the destination is not a user type, e.g. AbstractVector, but one of the other arguments is a user type? For example, I'd like to define something like:

struct Foo end
Base.broadcast!(f, dest::AbstractVector, A::Foo) = dest

but that conflicts with

broadcast!(f, dest::Union{SparseMatrixCSC, SparseVector}, mixedsrcargs::Vararg{Any,N}) where N in Base.SparseArrays.HigherOrderFns at sparse/higherorderfns.jl:1017

(specifically, my use case is re-implementing broadcast! in TypeSortedCollections with AbstractVector destination, current code here).

Third, if the non-destination arguments are Vararg, there is the additional issue that it's not possible to express the constraint that at least one of the arguments in a list needs to be of a certain type (see #21026 (comment) and following comments).

[timholy]

First, I'd like to request a bit more documentation on extending broadcast! for user types. In what follows I'm assuming that the idea is to simply add methods to broadcast!. Even if that's all, a sentence stating that that is the case would be appreciated.

That's about is, as far as I know. Once you get your broadcasting in TypeSortedCollections worked out, could I ask you to contribute that sentence?

but that conflicts with

Not for me:

julia> a = Foo()
Foo()

julia> broadcast!(sin, [1], a)
1-element Array{Int64,1}:
 1

You just mean that it causes an ambiguity? Ambiguities are not specific to broadcasting, and in this case I don't see a solution other than the usual one (if you want to nix the ambiguity, you need to define the method that resolves it).

Third, if the non-destination arguments are Vararg, there is the additional issue that it's not possible to express the constraint that at least one of the arguments in a list needs to be of a certain type (see #21026 (comment) and following comments).

Again, this is a general julia issue and not specific to broadcasting.

[tkoolen]

You just mean that it causes an ambiguity? Ambiguities are not specific to broadcasting, and in this case I don't see a solution other than the usual one (if you want to nix the ambiguity, you need to define the method that resolves it).

I can do that for this particular case, but it just seems very brittle, especially if a bunch of packages that may or may not be used together by users each extend broadcast!.

OK on the other points.

[timholy]

I can do that for this particular case, but it just seems very brittle, especially if a bunch of packages that may or may not be used together by users each extend broadcast!.

If you have an alternative, I'm all ears. How do you prevent ambiguities when some people are going to want to specialize on dest and other people are going to want to specialize on one or more of the inputs?

https://docs.julialang.org/en/latest/manual/methods/#man-method-design-ambiguities-1

[tkoolen]

One option could be to give the dest type precedence by dispatching on it first, as in https://docs.julialang.org/en/latest/manual/methods/#Dispatch-on-one-argument-at-a-time-1?

For example, Base could define

 broadcast!(f, dest::AbstractArray, As...) = broadcast_abstractarray!(f, dest, As...)
 broadcast!(f, dest::Union{SparseMatrixCSC, SparseVector}, As...) = broadcast_sparse!(f, dest, As...)

These appear to be the only specializations based on destination type currently in Base.

Both Base and user packages could then define methods for broadcast_abstractarray!, broadcast_sparse!, or broadcast_mytype! (together with a method broadcast!(f, dest::MyType, As...) = broadcast_mytype!(f, dest, As...)).

[timholy]

Probably a more scalable option would be (with apologies for the horrible name)

broadcast!(f, dest, As...) = broadcast_inputs!(f, dest, As...)
broadcast_inputs!(f, dest, As...) = # what the default broadcast! is now

Then if you're specializing on dest you extend broadcast!, otherwise you extend broadcast_inputs!.

Historically folks have not been happy with similar solutions (#10312). However, in practice it looks like we (unofficially) have this already: you could extend _broadcast! rather than broadcast!.

[tkoolen]

That's fine too. I can take a crack at it if people are happy enough with this approach. What are your thoughts on handling specializations based on the type of f?

[nalimilan]

These issues are really tricky. What would be the rationale for giving the priority to dest?

It seems to me that we should try to ensure that ambiguities do not make it impossible (by throwing errors) to combine array types which have not been designed to work together. Else, the broadcast machinery is going to be unusable in many cases.

Instead, we should fall back to the AbstractArray generic definition when no method has been defined for a given combination of types. Cannot this be achieved by using a mechanism similar to BroadcastStyle (or by reusing it)? People would either define a method dispatching on dest (as in @timholy's example above), or a method dispatching on the BroadcastStyle-like type computed from the inputs. The advantage of that approach is that promotion rules similar to BroadcastStyle can be used to choose which method is going to be called. It provides a fallback for cases which package authors did not explicitly anticipate, and allows stating things like "use this method if all inputs are either Array or MyCustomArray".

[Sacha0]

(Apologies for not finishing reviewing #23939 yet and being relatively quiet on this front! Working on things JuliaLang/LinearAlgebra.jl#57 is consuming all my bandwidth.)

Echoing Milan, one comment I planned to post in #23939 is that we should extend the wonderful mechanism therein to broadcast!, which has a similar need to compute which machinery a given call should be sent to. Best!

[timholy]

Agreed, this is an interesting design challenge, I like it 😈

Originally I didn't think we needed the BroadcastStyle machinery for broadcast!, as in my mind that's mostly about deciding on the correct kind of output container. However, TypeSortedCollections seems to be a good example of where it would be useful. Catching up to where @nalimilan and @Sacha0 are, I now think we should repeat the BroadcastStyle expansion for broadcast!. Something like:

broadcast!(f, dest, As...) = broadcast!(f, dest, combine_styles(As...), As...)
function broadcast!(f, dest, ::BroadcastStyle, As...)
    # "real" implementation
end

This seems open the door for specializing on the inputs while also limiting the potential for ambiguity by using the existing rules for combining inputs. In particular, this seems like a simple way to avoid any need to have separate variants depending on whether your tsc is the first, second, third, etc. input. This is essentially https://docs.julialang.org/en/latest/manual/methods/#man-methods-orthogonalize-1.

This change wouldn't force the same computation to be repeated twice, since broadcast(f, s::BroadcastStyle, ::Type{ElType}, inds, As...) can of course just pass s when it calls broadcast!.

However, this is only a small portion of the fix---there's still the question of how to handle specializations of 3 different arguments, f, dest, and the BroadcastStyle. By comparison with broadcast, I would hope that specializing on f would be extremely rare. (In the cases for broadcast there's a strong tendency to specialize one or more other arguments, too, which greatly reduces the risk of ambiguities.) In Base we only have the 3 identity specializations for broadcast!, and one of the 3 exists purely for ambiguity resolution. It's worth asking whether we should handle the other two inside the function body, e.g.,

           if isa(f, typeof(identity)) && N == 0
               if isa(A, Number)
                   return fill!(C, A)
               elseif isa(C, AbstractArray) && isa(A, AbstractArray) && Base.indices(C) == Base.indices(A)
                   return copy!(C, A)
               end
           end
          # The rest of this would be identical to current Broadcast._broadcast!

in which case Base would not have any specializations based on function type. (The if statements will resolve at compile time, so this shouldn't be a performance concern.) That doesn't provide a mechanism for packages to specialize on f, but like I said, I'm not certain we want to encourage that anyway.

With regards to how to specialize on dest or the BroadcastStyle: I think asking users to specialize broadcast_dest! is kind of ugly. For an overall cascade, what about:

broadcast!(f, dest, As...) = broadcast!(f, dest, combine_styles(As...), As...)
broadcast!(f, dest, ::BroadcastStyle, As...) = broadcast!(f, dest, nothing, As...)
@inline function broadcast!(f, C, ::Void, A, Bs::Vararg{Any,N}) where N
    if isa(f, typeof(identity)) && N == 0
        if isa(A, Number)
            return fill!(C, A)
        elseif isa(C, AbstractArray) && isa(A, AbstractArray) && Base.indices(C) == Base.indices(A)
            return copy!(C, A)
        end
    end
    # The rest of this would be identical to current Broadcast._broadcast!
end

Then the rule is:

• to specialize on dest, specialize broadcast!(f, dest::DestType, ::Void, As...)
• to specialize on the BroadcastStyle, specialize broadcast!(f, dest, ::BroadcastStyle, As...)

This explicitly gives precedence to the BroadcastStyle. It follows the strategy of https://docs.julialang.org/en/latest/manual/methods/#Complex-method-%22cascades%22-with-default-arguments-1

[tkoolen]

I like @timholy's suggestion, definitely the BroadcastStyle part. I was kind of expecting @nalimilan's comment, but decided not to bring it up myself as I was not too familiar with the BroadcastStyle code yet. This also addresses my third point in the issue description, correct?

After thinking about it for a bit, I like the handling of f as well. No need to over-engineer this in Base if there are only so few cases to handle, and I think it leaves enough room for the rare case that user packages do need to specialize on f as an optimization (for a given list of dest types or BroadcastStyles). The Base.indices part of the f 'specializations' would of course not be resolved at compile time, but that's inside the identity case anyway and would result in the same code as with the current implementation.

[nalimilan]

Great proposal!

I'm not completely sure we'll be able to get rid of all specializations on f, but I don't see how it could be handled better anyway. We would need to have a more precise idea of when specializing on f is needed to find out what kind of mechanism could be useful. Anyway your proposal isn't worse than the current situation in that regard, and it sounds reasonable enough to handle identity as a special case.

[vtjnash]

That all sounds fairly complicated. What do you think about doing user-controlled lazy-fusion instead? I started a branch to demonstrate that, I just stopped short of finishing updating some of the tests. If anyone wants to take it over, it's at https://github.com/JuliaLang/julia/compare/vtjnash:jn/lazydotfuse.

[nalimilan]

Can you explain how it would help with the problem at hand?