Julia: Element-wise call syntax does not always lower to standard broadcast()

[nalimilan]

It appears that some simple uses of the element-wise call syntax .( are not lowered to a standard broadcast call. This is a problem for NullableArray where we would like to overload broadcast to apply lifting semantics (https://github.com/JuliaStats/NullableArrays.jl/pull/166).

For example, the following code doesn't call broadcast(get, X, 0). Therefore it fails since my custom broadcast function has a special behavior for get which doesn't get triggered here.

Why it this the case?

julia> expand(:(get.([Nullable()], 0)))
:($(Expr(:thunk, CodeInfo(:(begin 
        $(Expr(:thunk, CodeInfo(:(begin 
        global ##3#4
        const ##3#4
        $(Expr(:composite_type, Symbol("##3#4"), :((Core.svec)()), :((Core.svec)()), :(Core.Function), :((Core.svec)()), false, 0))
        return
    end))))
        $(Expr(:method, false, :((Core.svec)((Core.svec)(##3#4,Any),(Core.svec)())), CodeInfo(:(begin 
        return get(#temp#,0)
    end)), false))
        #3 = $(Expr(:new, Symbol("##3#4")))
        SSAValue(0) = #3
        SSAValue(1) = (Base.vect)(Nullable())
        return (Base.broadcast)(SSAValue(0),SSAValue(1))
    end)))))

[yuyichao]

This is an optimization since broadcast(get, X, 0) shouldn't do anything other than what the special lowering does.

[nalimilan]

So that means one cannot specialize broadcast for a custom array type? Sounds quite restrictive to me. Is the expected behavior documented somewhere?

[nalimilan]

To clarify, IIUC, get.(X, 0) gets lowered to broadcast(x -> get(x, 0), X) instead of broadcast(get, X, 0) for performance reasons. Since that's done at parse time, no dispatch on the input types is possible. I see the point of this optimization, but unfortunately it hardcodes the behavior of broadcast for any type in the parse, which is unusual for a Julia method which by definition can be overridden.

Maybe there should be a type check, and the fast code would only be used for Array and Base scalar types?

[martinholters]

If you need to specialize broadcast for get, that is a bad sign anyway since it wouldn't work with fusion. E.g. neither get.(foo.(X), 0) nor get.(foo.(X), i) would end up calling your specialization.

[nalimilan]

Yes, I've noticed that too. Whatever we try to improve the Nullable workflow, we bump against lack of support in Base. I'm starting to get tired of these repeated dead ends. See https://github.com/JuliaStats/NullableArrays.jl/pull/166.

[stevengj]

Yes, as I commented in #17623, the existence of parse-time fusion by itself (even without constant-folding optimizations) makes defining specialized broadcast methods for particular functions effectively useless.

On the plus side, this sometimes forces us to pursue more general solutions, e.g. for preserving sparsity for broadcast on sparse matrices.

[nalimilan]

Right. Unfortunately, we can't just call the function to check whether it is zero-preserving or looks at its return type to decide whether to return a BitArray. Here, non-lifting functions really need to be identified and transformed before being fused with others. The generalization that this encourages is to lift all function calls involving Nullable by default, even outside broadcast.

[stevengj]

@nalimilan, what is a "non-lifting" function and how would you identify one at parse-time?

[nalimilan]

I don't think hard-coding them at parse time would be a good idea, even though their list should be pretty limited. Basically, these are functions which operate on the Nullable object itself rather than the value it wraps: isnull, get, plus logical operators isequal, isless, & and | (for three-valued logic). See JuliaStats/NullableArrays.jl#166 and https://github.com/JuliaStats/NullableArrays.jl/issues/144.

[stevengj]

I guess in principle we could transform:

f.(g.(x...)) ---> isfusing(g,x...) ? broadcast((_...) -> f(g(_...)), x...) : broadcast(f, broadcast(g, x...))

where isfusing would default to true. i.e. you allow the caller to define "fusion barriers" for particular functions. Or isfusing could automatically return false if a specialized broadcast method exists for g? I'm not sure if this is a good direction to pursue, however.

[stevengj]

Note that if you want a non-fusing broadcast-like operation, you can just define a method without dots that acts on arrays. e.g. & and | will never be fusing. Similarly, you could just define get(a::NullableArray, ...) that acts like a get.(a, ...) but doesn't fuse.

[nalimilan]

That would work. Though strictly speaking loops can fuse, it's just that we need to be able to transform the function if needed before fusing it. That is, something like:

f.(g.(x...)) --> broadcast(_ -> maybe_lift(f, maybe_lift(g, _...)), x...)

with

maybe_lift(f, _...) = f(_...)

for non-lifted functions (the current case).

Of course we can always define special methods instead of using the dot syntax. But I prefer trying to push the logic as far as possible to see all implications and (in)consistency of the design choice first.

[stevengj]

But if you do e.g. sqrt.(abs.(x)) on a nullable array, wouldn't you want to lift only once, not once for abs and once for sqrt?

[nalimilan]

Ideally, yes. If we inline the lifting logic, the compiler could be able to do the checks only once (haven't checked). If not, a smarter design would be:

f.(g.(x...)) --> broadcast(_ -> maybe_lift((f, g), x...))

maybe_lift would lift only once for all subsequent calls that need it.

[stevengj]

@nalimilan, that sounds like it could be implemented without any parser changes at all. Couldn't you just define something like

Base.broadcast(f, ...arguments including a NullableArray...) =
    invoke(broadcast, (...NullableArray replaced by AbstractArray...), (args...) -> lift(f, args..), ...arguments...)

[nalimilan]

That's what I'm doing currently. It works when all functions lift, but it fails when fusing a non-lifting function like isnull with another one. I really need access to the original function before fusing to know what to do.

[JeffBezanson]

Trying to automatically tell whether a function should be lifted or whether it expects to see Nullables is never going to work. It's not even a well-formed question; for example is tuple a lifting function? You can make tuples of anything, so there's no way to know.

[nalimilan]

@JeffBezanson There's nothing automatic in the mechanism I described. Only functions for which lift has been overridden would use a behaviour different from the default (which is lifting). Is that an issue?

[davidagold]

To be clear, the list of such functions is quite short: &, | (for three-valued logic), isnull, get. Perhaps unsafe_get belongs in there, too. It's possible that this list expands, but if it does, it hopefully does so very, very slowly.

EDIT: Technically, & and | are not "non-lifting" -- they still need to be lifted for three-valued logic semantics, but we define specialized lift(::typeof(&), x, y) methods.

[nalimilan]

I've made some investigations regarding a concrete implementation. It seems that we could replace the current function fusing code in the parser with a call to broadcast_fuse(types::Type, f, fs...) which (by default) returns an anonymous function composing f and fs.... The types argument would be used to pass the types of the inputs to broadcast, which would allow providing a custom method for NullableArray.

Performance-wise, this gives the same results as the current approach AFAICT:
https://gist.github.com/nalimilan/e97e6561c7006b4ce2dba0f85519a2a8

@stevengj Do you think it would be possible/not too hard to change the parser code to do that? This could potentially allow simplifying that code too. (Unfortunately, my Scheme skills are close to null so I'm having a hard time trying to assess this...)

[martinholters]

IIUC, the fusing is a bit more complicated, as it needs to account for functions with more than one argument. Given only f and g in broadcast_fuse, how would you discern f.(g.(a,b,c)), f.(g.(a,b),c), f.(a, g.(b,c)), f.(g.(a),b,c), ...?

[nalimilan]

Good catch. So I need to find a more complex approach.

[johnmyleswhite]

I think part of this conversation is complicated by the use of the term "non-lifting" (which I would tend to also call lifting) and distinctions between parse-time and later periods when indirection might change behaviors.

As I understand Milan's goal, what he really wants is the following:

• Broadcast is syntactically replaced in two parts rather than the current one part:
  
  
  
  • First, replace every single call site of f with lift(f)
  
  
  • Then do everything else we're currently doing
  
  
• Then we need to assume the following:
  
  
  
  • lift(f) generates a new callable that, given non-nullable typed arguments, behaves identically to f.
    
    
    
    
    
    • If there are any nullable typed-arguments, lift(f) will, by default, extend the behavior of f from its non-nullable typed definition so that it generates normal outputs if all arguments are non-null values, but generates a null value if any arguments are a null value.
    
    
    
    • If there are any nullable typed-arguments and lift(f) has had custom methods added to it, then those custom methods will allow the behaviors that Milan has been calling "non-lifting".
    
    
    
  
  

So basically everything would be lifted, but calling the lifted function with non-nullable arguments is identical to calling the original function. Only in other cases (where arguments are Nullable, NullableArray, etc.) would this additional lift(f) indirection lead to behavior different from working with f all the way through.

[nalimilan]

Indeed, a simpler approach would be to always add lift calls from the parser, since they are no-ops when no argument is a Nullable.

[johnmyleswhite]

Assuming we don't hit any performance issues dealing with all the inlining that is likely to be needed, would there be objections to such an approach?

[stevengj]

It seems a little weird to add a parser transformation that seems like it would only be usable by a single package...?

[johnmyleswhite]

I think this approach is part of our attempts to propose making nullables a much larger part of Julia's design (and to move all of the basic lifting into Base Julia). Partly this came up when discussing C#'s design with Erik Meijer and Eric Lippert, who had originally proposed a design for C# that is very similar to what I think Milan would like to see in Julia.

[martinholters]

How much of a headache is it to write lift given that any function applications in it would also be rewritten? Would we need special syntax for function application without implicit lift?

[nalimilan]

@martinholters I'm not sure I understand your question. Can you develop?

[martinholters]

Say, for the default case, lift constructs a Lifted{F<:Function} where (::Lifted)(args...) has the required overloads:

immutable Lifted{F<:Function}
    f::F
end
(l::Lifted)(args...) = l.f(args...)
# and overloads for Nullable arguments...
lift(f) = Lifted(f)

If the parser turns this into

immutable Lifted{F<:Function}
    f::F
end
(l::Lifted)(args...) = lift(l.f)(args...)
# and overloads for Nullable arguments...
lift(f) = lift(Lifted)(f)

I see infinite recursion looming...

[nalimilan]

No, the parser would only transform broadcast calls (cf. the OP).

[martinholters]

Ahhh, makes sense, sorry for the noise.