conservative constant prop if broadcasting is involved

[antoine-levitt]

Here's my best attempt at reproduction (the actual case is https://github.com/JuliaMolSim/DFTK.jl/blob/master/src/fft.jl#L47)

function f(x; isreal=true)
    y = similar(x)
    y .= x
    isreal ? real(y) : y
end
@code_warntype f(randn(ComplexF64, 1))

gives a ::Union{Vector{ComplexF64}, Vector{Float64}}. This is on 1.7 and on master. Removing the y .= x or moving isreal to a Val both fix the issue. y[:] = x[:] also has the inference failure.

[martinholters]

Can be fixed with Base.@constprop :aggressive. Alternatively, you can keep the method body complexity low enough by explicitly outlining the broadcast:

function f(x; isreal=true)
    y = similar(x)
    @noinline y .= x
    isreal ? real(y) : y
end

(Obviously, this requires sufficiently recent Julia. Otherwise, the broadcast has to be moved to a helper function marked @noinline.)

That said, not having to do anything special here and making things just work would certainly be desirable.

[aviatesk]

I'm not sure if we want to enable constant propagation on this function though -- it involves a broadcasting and so a fair amount of complexity from the compiler view point.
In other word, we won't inline f anyway and so there is not much benefit in doing constant propagation on f in terms of inlineability (though there is certainly some profitability in terms of type accuracy).

For now, I'd suggest encoding some information that is equivalent to the isreal condition in the type signature of f (probably still better to avoid using Val), and not rely on constant-propagation to obtain the desired type accuracy.

[antoine-levitt]

Thank you! I was actually not aware that constant-propagation was based on the body complexity, I thought it was based on the type of the arguments, and booleans propagated. Val is a good enough workaround.

[antoine-levitt]

It'd still be nice if it worked automatically. Constant-propagating literals (booleans at least, maybe ints and floats?) looks like a good idea in general, but there might be a reason I don't see why it's not done.

[aviatesk]

The heuristic to decide whether or not we do constant-prop' really influences the compilation latency. It would be better to come up with some better heuristic to enable constant prop' in this kind of situation (i.e Bool-lean call argument used for branching sometimes directly influences return type inference) since it can give us a better type information, though I don't think it can be easily implemented.

And this pattern of code (i.e. whose return type depends on value of a call argument) should be avoided if you want to get best performance, as noted in performance tips.

[antoine-levitt]

OK, thanks for the response! I understand finding good heuristics that work reasonably well for a lot of different use cases is super hard.

And this pattern of code (i.e. whose return type depends on value of a call argument) should be avoided if you want to get best performance, as noted in performance tips.

I know, of course, but If I wanted that I'd use Fortran ;-) Type-unstable code is much more flexible for designing APIs. The whole point of julia for me is that you can use this kind of patterns in functions that are not performance critical, and just have make sure your performance-critical code is insulated from it. it's just simpler when the compiler figures it out for me and I don't have to work for it.