Why are there no captured variables when the `if` expression is used as an rvalue?

New to Julia
, ,
1

When I was learning the Performance of captured variable, I found a second version of the example function that shows no captured variable in @code_warntype. I am confused about the following three questions:

  1. Would these two functions logically equivalent?
  2. Would the anonymous functions produced by these two functions behave the same way?
  3. Why does the first function exhibit variable capturing while the second function does not, despite using different if expressions?

Thanks for any responses.

Code details 
function abmult(r::Int)
    if r < 0
        r = -r
    end
    f = x -> x * r
    return f
end

function abmult_new(r::Int)
    r = if r < 0
        -r
    else
        r
    end
    f = x -> x * r
    return f
end
4 Likes
2

You should expect the same results from the returned functions, but they won’t perform the same. In both cases there’s capturing of the variable r by x -> x * r. What’s different is abmult boxes the captured variable, that’s the Core.Box you’re seeing in @code_warntype. A box provides semantically correct flexibility if the variable can be reassigned to an instance of a different type after capturing. However, r can not be reassigned after capturing, and we ourselves can infer the stable type of r throughout. Turns out abmult is doing this unnecessarily.

Digging further, each nested function definition gets 1 underlying type like typical named functions; we wouldn’t want ballooning overhead of types and their method tables. You can see a bit of the definition’s structure in @code_warntype, but it’s possible to display it more obviously:

julia> dump(typeof(abmult(1)).name.wrapper)
var"#3#4" <: Function
  r::Core.Box

julia> dump(typeof(abmult_new(1)).name.wrapper)
UnionAll
  var: TypeVar
    name: Symbol r
    lb: Union{}
    ub: Any
  body: var"#1#2"{r} <: Function
    r::r

abmult’s type captures the r variable in a boxed field (kind of like a Base.RefValue{Any}), while abmult_new’s type captures r in a parametric field, which allows inferring a stable type for r. abmult_new’s parameter looks like an obvious win for type inference…at first.

In the general case, call-wise type inference is unfeasible for variables shared by multiple methods. When the first outer method is called and compiled, you need to infer the variable’s type but you can’t narrow it down so much that the other methods fail to compile on their calls. So, we fall back to an uncommitted Core.Box, but we can try our best to find optimization opportunities. For example, if none of the nested methods ever reassign the variable, we could infer the variable’s type from the first outer method’s call and set that as the parameter. Unfortunately, we currently make the boxing decision when the method definitions are lowered, far before any call or compilation where type inference could help. It roughly and conservatively considers whether the variable is reassigned in a parameter-worthy way.

The absence of type inference in determining the parameter means abmult_new’s decision actually isn’t a clear win for type stability. Consider this type-unstable version:

julia> function abmult_newer(r::Int)
           r = if r < 0
               -r
           else
               Ref(r)
           end
               f = x -> x * r
           return f
       end
abmult_newer (generic function with 1 method)

julia> abmult_newer(2) |> dump
#17 (function of type var"#17#18"{Base.RefValue{Int64}})
  r: Base.RefValue{Int64}
    x: Int64 2

julia> abmult_newer(-2) |> dump
#17 (function of type var"#17#18"{Int64})
  r: Int64 2

You could see this a bit better with the color coding of @code_warntype, but now the anonymous function is type-unstable, using the runtime input’s type at instantiation like Ref calls do. Although inferring two types for the parametric type is arguably equally bad as inferring two types for the parameter instead, @code_warntype currently doesn’t even infer the anonymous function’s type as these two types, despite inferring r itself as a union of two types.

For now, if we want that type inference optimization given no reassignments in nested methods or after them, we need the let trick to make a new assigned-once variable for capturing. For better control of the captured types, we could make function-like objects and instantiate them in place of nested methods.

2 Likes
3

To address the why: I think the key difference is that in your version, the variable is always reassigned at one particular point in the function body. In such cases, the compiler can effectively treat r before reassignment and r after reassignment as two different variables. It’s as if you wrote

function abmult_new(r::Int)
    s = if r < 0
        -r
    else
        r
    end
    f = x -> x * s
    return f
end

Since s is of known type never reassigned, boxing is avoided.

I’m not the right person to explain the finer points and lower-level details, but I asked about this on slack once and got an answer that used the term phi node, so that could be a starting point for further googling.

4 Likes
4

This is visible as a duplication between @code_warntype’s division of arguments and local variables, and @code_warntype abmult_newer(2) shows off the divergence in type inference.

It’s not clear what the internal criteria are, but this lines up with what I’ve seen and is put the best way I’ve heard so far.

2 Likes
5

One of the most common uses of this capability (and the one I asked about in that slack post) is the following pattern:

function f(x, y=nothing)
    if isnothing(y)
        y = default_y(x)
    end
    # actual function body
end

Since isnothing(y) is resolved at compile time, the reassignment of y = nothing to y = default_y(x) either always happens or never happens in any given methodinstance, and consequently, the “edition” of y used in the remainder of the function body will be inferred as Ytype rather than Union{Nothing,Ytype}, even though the variable y as defined by the scoping rules has both types over the course of execution.

(I know it would be trivial to avoid ever setting y = nothing in this MWE, but I’ve had some situations where the only real option was to put a default of nothing for an optional argument, even though an actual default value would always be computed and assigned at the top of the function body.

2 Likes
6

Also a good way to highlight that the boxing decided at lowering can ruin otherwise good type inference at later compilation:

julia> function fclosure(x, y=nothing)
           if isnothing(y)
               y = x
           end
           y, (() -> y)
       end;

julia> @code_warntype fclosure(1, nothing)
MethodInstance for fclosure(::Int64, ::Nothing)
  from fclosure(x, y) @ Main REPL[21]:1
Arguments
  #self#::Core.Const(Main.fclosure)
  x::Int64
  y@_3::Core.Const(nothing)
Locals
  #13::var"#13#14"
  y@_5::Union{}
  y@_6::Union{}
  y@_7::Union{Nothing, Core.Box}
Body::Tuple{Any, var"#13#14"}
...

If y weren’t captured, it’d be inferred as typeof(x) when returned. As said before, it’s hypothetically possible because the closure never reassigns y. Actually, even if its body reassigned y, we would still be able to infer the returned y value’s type because the closure is never called within the scope where it is assigned to y. However, that’s a difficult thing to determine in general because a reassigning closure could be passed to a (possibly uninferred) higher order function where it is called, so we could at most conservatively check whether a closure isn’t called or passed as an argument in the scope. The added work to compilation has to be mitigated or justified, too.

2 Likes