Can the type gymnastics in broadcast.jl be replaced with @generated?

I am trying to learn from base/broadcast.jl and eventually implement a custom broadcaststyle, which is a bit more involved than can be implemented by the methods in doc. I see the following type gymnastics and am wondering whether I can use a simple @generated function.

I am wondering what is the downside of using:

@generated function BroadcastStyle(a::A, ::DefaultArrayStyle{N}) where {A<:AbstractArrayStyle{M}} where{M,N}
    if M>=N
        return :(A(M))
    else
        return :(A(N))
    end
end

instead of:

DefaultArrayStyle(::Val{N}) where N = DefaultArrayStyle{N}()
DefaultArrayStyle{M}(::Val{N}) where {N,M} = DefaultArrayStyle{N}()

BroadcastStyle(a::AbstractArrayStyle{M}, ::DefaultArrayStyle{N}) where {M,N} =
    typeof(a)(_max(Val(M),Val(N)))

_max(V1::Val{Any}, V2::Val{Any}) = Val(Any)
_max(V1::Val{Any}, V2::Val{N}) where N = Val(Any)
_max(V1::Val{N}, V2::Val{Any}) where N = Val(Any)
_max(V1::Val, V2::Val) = __max(longest(ntuple(identity, V1), ntuple(identity, V2)))
__max(::NTuple{N,Bool}) where N = Val(N)
longest(t1::Tuple, t2::Tuple) = (true, longest(Base.tail(t1), Base.tail(t2))...)
longest(::Tuple{}, t2::Tuple) = (true, longest((), Base.tail(t2))...)
longest(t1::Tuple, ::Tuple{}) = (true, longest(Base.tail(t1), ())...)
longest(::Tuple{}, ::Tuple{}) = ()

Any help appreciated.

Edit:
After consulting master, the issue is already addressed by this commit. That seems like a miracle to me, thank you Julia team!

However I still see this :

Base.@pure function BroadcastStyle(a::A, b::B) where {A<:AbstractArrayStyle{M},B<:AbstractArrayStyle{N}} where {M,N}
    if Base.typename(A) === Base.typename(B)
        return A(Val(max(M, N)))
    end
    return Unknown()
end

Can you explain why @pure is necessary?

A disadvantage of generated functions is that they are really two functions in one: the generator (what runs at compile time) and the code that gets generated (what runs at runtime). As a consequence, rules governing dispatch, backedge-triggered recompilation, etc, are either more complex, not as mature, or simply do not apply. For example, suppose you call a generated function, thus triggering the generator, and then later add a more specific method for a call made by the generator: should you regenerate the code? Probably yes, but right now that doesn’t happen. As a consequence, you can’t update generated functions automatically with Revise.jl, etc.

As a consequence, for some of the things that would be easy to do with generated functions, there is some incentive to do them by other ways. Once you learn lispy tuple programming it’s not so hard, even if it results sometimes in complicated-seeming dispatch hierarchies. Fortunately, better constant-propagation has meant that we can often get away with simpler constructs.

Marking the method as @pure is an intermediate solution: for the purposes of inference it mixes runtime and compile-time, and can help fix inference problems in cases where the existing inference mechanisms do not suffice on their own. The most obvious case of needing to mark something as @pure would be for a ccall that computes something involving types, since the result of the ccall is unobservable by inference. However, this annotation can be useful in other cases too, such as the one you noticed (because of the typename part which is not fully inferrable).

Incorrectly marking something as @pure can lead to bugs if you get it wrong, leading to a running joke among Julia developers (that also links to a case where the lack of a @pure annotation led to a problem). But the general rule is that if it operates purely in the type domain and has no side effects, it should be fine.

10 Likes

That’s super enlightening. I must say everytime I get an answer from you it has always been a pleasure to read, and a wealth of material to learn from. Thank you and Merry Christmas!

2 Likes