Can someone explain closures to me

Internals & Design
, ,
1

Hi, I am familiar with what a closure is (a function that captures state outside its scope) but I am interested in how they work in Julia so that I can understand performance.

In Performance Tips there is a section dedicated to closures (Performance of captured variable).

Namely it explains that the following is poor performance:

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

And the next is better:

function abmult2(r0::Int)
    r::Int = r0
    if r < 0
        r = -r
    end
    f = x -> x * r
    return f
end

And lastly the below is best:

function abmult3(r::Int)
    if r < 0
        r = -r
    end
    f = let r = r
            x -> x * r
    end
    return f
end

I do not see how any more information is added in any of the examples (i.e if I were a compiler I would compile them all the same). It is clear to me that r is an Int when it is captured. It is also clear to me that no other piece of code can change r. So why is it not clear to the compiler?

Basically I think without a technical explanation of some internals I can’t understand what is going on in this example or any use of functions as values in general (not just functions that capture state outside their scope, a.k.a closures). Another thing I don’t understand is how a function that looks like what is below is compiled:

call_function(f, x) = begin
    f(x)
end

Is a new specialization for call_function created on every invocation with a different f?

For example:

call_alot_of_functions(xs) = begin
    squares = similar(xs)
    for (i, x) in enumerate(xs)
        f = y -> y * x
        squares[i] = call_function(f, x)
    end
    squares
end

Is this going to be horribly slow? Is f going to be of a different type on each iteration of the loop? Leading to call_function having to be specialized on each iteration? And is f itself going to be recompiled each time?

Surely using closures and anonymous functions isn’t super unperformant given their prevalent use, but I just can’t wrap my mind about how they work internally and feel like it is a hole in my understanding of how the compiler will treat the code that I write

7 Likes
2

This longstanding issue has moved a bit since the beginning but is far from over, hence the tip. I’m not certain on the exact boundaries of captured variable inferrability, but here’s what I know so far, starting from some basics.

A closure is a function that uses an outer local scope’s variables; this is not different from how other local scopes reuse outer local variables. The distinction of “capturing” is that a nested function is also an instance that can outlive the outer scope, so captured variables can also outlive the scope. It’s important to note that as a consequence, if the captured variable is reassigned anywhere, that change is also shared by all closures.

Example of reassignment of captured variables. 
julia> let
         v1 = 0
         f() = v1
         println(f())
         v1 += 1
         g() = v1
         println((f(), g()))
         v1 += 1
         println((f(), g()))
       end
0
(1, 1)
(2, 2)

A method is compiled separately for each unique call signature, which is the types of the function and the ordered runtime arguments. The compiler uses the argument types to infer the types of values and local variables, so inferred types vary among call signatures. Conversely, a variable captured from an outer scope must have the same inferred type in that scope and in the infinite possible call signatures. A captured variable cannot be inferred if it is reassigned in a closure because the compiler cannot infer the value across the infinite call signatures it has not encountered yet.

A reassigned captured variable has an unfixed type. 
let
  x = 0 # x::Int
  function setx(y)
    x = y # setx captures x, shares x with let block
  end
  setx(1.0) # x::Float64
  setx(true) # x::Bool
  setx(1im) # x::Complex{Int}
  # ... so the compiler can only say x::Any
end

Since a call-wise compiler can’t be exhaustive, the parser lowerer (or whatever the noun is) checks the expression, and if a variable is never reassigned or has a type declaration, the type is known to be unchanged, despite only being a symbol at that point.

Closures are implemented as functors; each captured variable resides in a field of a hidden struct’s callable instance. Uninferrable or reassigned variables are stored in ::Core.Box fields, which is like an internal version of Ref{Any}; this includes type-declared variables like r::Int in abmult2, the type being restored by convert and assert steps. Inferred and non-reassigned variables are stored in parametric fields ::T; you can see this parameter show up in closure’s types. Given the inherent difficulty of inferring variables shared across call signatures, I think the biggest possible improvements would be 1) doing ::T when all reassignments occur prior to all closures, like in abmult, 2) doing ::Ref{T} when the shared variable has 1 type declaration (>1 causes a syntax error) and is reassigned within or after any closure, and 3) inferring closures better when a captured variable is inferred with a small Union type; someone correct me if there are any misconceptions there.

Inferred and non-reassigned variables contribute type parameters to closures. 
julia> let
         x = 1
         f() = x
       end
(::var"#f#1"{Int64}) (generic function with 1 method)

julia> let
         x = 1
         x = 2 # x is no longer inferrable
         f() = x
       end
(::var"#f#2") (generic function with 1 method)

Inferring a variable doesn’t always make things type-stable though. Just today, there was a thread demonstrating that inferring a variable as T::DataType doesn’t make for inferrable calls e.g. one(T). The most adaptable workaround was refactoring to capture an instance of T.

15 Likes
3

The problem with closures arises when you have a function defined inside a function AND you’re redefining variables/arguments. The compiler gets confused and you end up with type unstable code.

For example, note in this example everything is fine when you don’t redefine a.

function outer_foo()
    a = 1
    a = 1
    inner_foo() = a 
    return inner_foo()
end

@code_warntype outer_foo() # type unstable


function outer_foo()
    a = 1
    b = 1
    inner_foo() = a
    return inner_foo()
end

@code_warntype outer_foo() # type stable

I recommend not using closures if possible. Most of the time, you can redefine the inner function outside the outer function.

If you really need to use closures, you should be fine in most cases when you don’t redefine variables or function arguments. But always check with @code_warntype.

1 Like
4

I only halfway agree because closures are often a lot more convenient to write than functors and methods, especially with comprehensions and generators that are built on closures. If performance isn’t important, it’s fine; if performance is important, it’s fine if variables are not reassigned.

However, explicit functors are nice because it’s easier to reason about stored instances rather than captured variables. A shared variable can be type-stably mimicked by passing an inferrable instance as an argument to the functors’ type constructors. The compiler can infer this much better because no actual variables are shared by multiple methods. A reassignable variable can be mimicked by storing a Ref{T} instance.

1 Like
5 
begin
    a = 1
    f() = a
    g() = a
    f, g
end

So in this case can the parser tell that a is not reassigned and do proper type inference? Or does a variable being captured by multiple closures cause inference to result in Ref{Any} for the field of the function object? Also I guess is the “proper” thing to do here Ref{Int} or just Int

6

That is different from what I described earlier because f and g are functions in the global scope accessing a global variable a; unlike local scopes, global scopes persist and cannot be outlived, so my understanding is that they are not considered closures in Julia because they do not need to capture anything in fields. In any case, a is an global variable with no type or const declaration, and those can only be inferred as ::Any because they can be reassigned arbitrarily.

If it were a local let block, you would be correct that a is inferrable.

This does not happen, Julia uses Core.Box, which has more differences than just the type and name; the relevant similarity is that the type of the contained instance is unfixed.

7

I love coding with closures because it makes it so easy to reuse memory buffers and temporary variables without explicitly storing them.

Such as:

function build_forward(A, B)
    f = let buffer = similar(A), A=A, B=B
        function f(x)
             buffer .= # do something with x, A, B
             buffer .= # do something with x and A
             return B
         end
   end
   return f
end

But yeah, with the let ... end it’s always quite some boilerplate.

8

It may be helpful to note that closures of re-assigned variables are a really weird language construct.

Normally, variables in julia are bindings; mere shorthands for the value that was last assigned to them. They are not real objects that exist in any way:

function foo()
   r_contents = 1
...

If you want to reify the variable into a kind of “slot” that can be written and read, you can do that explicitly:

function foo()
   r_address = Ref(1)
...

This is similar to C:

void foo(){
   long r_contents = 1;
...
}

void foo(){
   long r_address[1];
   r_address[0] = 1;
...
}

Of course the compiler is permitted to look at your code and try to de-reify variables, i.e. try to replace the r = Ref(1) by r_contents = 1, and then replace all uses r[] by r_contents . In llvm, that’s the job of mem2reg / memssa passes.

Now C has a really weird language construct: The addressof operator. Instead of having to declare types, you can do

void foo(){
   long r_contents = 1;
   long* r_address = &r_contents;
...
}

This is pretty surprising: The C compiler is expected to do type inference in order to figure out whether a variable is a binding, or a reified slot. And this type inference doesn’t follow control-flow, it’s a non-local syntactic property of the code:

void foo(){
   long r_contents = 1; // looks like a binding
   if(false){
      // this code is dead. Doesn't matter, it's a syntactic nonlocal property of your code
      // let's promote r_contents from binding to reified slot, after the fact!
      long* r_address = &r_contents;
  }
...
}

Luckily julia doesn’t do such nonsense. You must decide from the beginning whether your variable is supposed to be a binding or a reified slot. And this is a local property, so it’s easy to understand code when reading it.

…just joking. Julia does not have an explicit addressof operator (imo for very good reasons!).

But closures work this way:

function foo()
    r_contents = 0; # looks like a binding
...
   return
   # this code is dead. Doesn't matter, it's a syntactic nonlocal property of your code
   # let's promote r_contents from binding to reified slot, after the fact!
    get_r() = r_contents
    set_r(v) = (r_contents = v;)
end

This sucks for the compiler and the reader of your code :frowning:

edit: So the punchline regarding your original question is:

Closures in julia do two things:

  1. they capture state from outside their scope. This is not an issue and has no surprising perf impact. Like a functor.
  2. They don’t just create a functor, they implicitly promote the captured state from outside their scope from binding to addressable object – and this is syntactically invisible “outside their scope”. This is where all the surprising perf impacts come from.
2 Likes
9

My advice to not use closures was more along the lines of “unless you know what you’re doing, avoid it if you can”.

I agree that closures arise naturally and make the code easier to read. But they can easily make your code type unstable even when you’re writing code in a typical Julian way.

I was simplifying also when I said that the problem is when you reassign variables. Consider the following.

Example 1:

function foo()
    bar() = x
    x     = 1       

    return bar()
end

@code_warntype foo()      # type unstable

In fact, the following doesn’t solve the problem:

function foo()
    bar()::Int64 = x::Int64
    x::Int64     = 1       

    return bar()
end

@code_warntype foo()      # type unstable

but this does

function foo()
    x            = 1
    bar()        = x
    
    return bar()
end

@code_warntype foo()      # type stable

so that the order in which you include the functions would matter.

Example 2:

function foo(x)
    closure1(x) = x
    closure2(x) = closure1(x)
    
    return closure2(x)
end

@code_warntype foo(1)            # type stable

function foo(x)
    closure2(x) = closure1(x)
    closure1(x) = x
    
    return closure2(x)
end

@code_warntype foo(1)            # type unstable

In both cases, you can fix it by making the inner function have everything that it’s used as an argument. But this would include functions. So, if you don’t want to care about the order of functions, you’d have to do the following:

function foo(x)
    closure2(x, closure1) = closure1(x)
    closure1(x)           = x
    
    return closure2(x, closure1)
end

@code_warntype foo(1)            # type stable

Overall, every time I use closures I double check that I didn’t add some type instability without noticing it.

1 Like
10

It doesn’t appear that you need the let block there, function build_forward already makes a local scope that holds a local A and B, and it could hold a local buffer too. You only need that pattern when you want to reuse variable names without the closures actually sharing variables.

The functor seems to store x in a Core.Box because x is unassigned at that point. Unlike a global function that can wait until call-compilation to figure out global variables, a nested function definition immediately instantiates a functor whose fields are set in stone to contain the captured variables. Maybe the parser could also improve to recognize that no closures are called before the captured variable is assigned once and defer the functor instantiation? Not sure, could probably come up with pathological situations.

I don’t see that, I see the typical effect of annotating x::Int = 1 when it is assigned after a closure. Annotating bar() was unnecessary.

`@code_warntype` report. 
julia> @code_warntype foo()
MethodInstance for foo()
  from foo() in Main at REPL[1]:1
Arguments
  #self#::Core.Const(foo)
Locals
  x::Core.Box
  bar::var"#bar#1"
Body::Int64
1 ─      (x = Core.Box())
│        (bar = %new(Main.:(var"#bar#1"), x))
│   %3 = x::Core.Box
│   %4 = Base.convert(Main.Int64, 1)::Core.Const(1)
│   %5 = Core.typeassert(%4, Main.Int64)::Core.Const(1)
│        Core.setfield!(%3, :contents, %5)
│   %7 = (bar)()::Int64
└──      return %7

Again, an unassigned closure1 was stored in a Core.Box.

On a wider topic, it’s expected that nested functions, including closures, are more limited than global functions. In a local scope, you shouldn’t (and have no good reason to) overwrite method signatures, nor can you feasibly conditionally define methods; those linked issues are related, and the limitations are necessary for letting those underlying functors be performant. Hopefully one day those limitations can be laid out clearly in the docs and in errors, maybe loosened a little with an improved parser-lowerer. Having written this out, I’m agreeing with this a bit more:

If there’s no time to understand local scoping rules and still-obscure nested function limitations, use functors to store state for calls, they’re a great and underappreciated core feature (functions are built on stateless functors). Fields holding mutable instances like Ref{T} allow functors to share state.

11

Thanks for the pro tip. This erased most of my concerns about closures in Julia. Lack of ergonomic and performant closures would be very unsatisfying if you want to write programs in a functional style, but treating variables as bindings (that are never reassigned) is exactly what’s encouraged by a functional style, so this problem solves itself.

1 Like
12

I agree that functors are great. But they can be frustrating to use if they are intended for some one-off calculation. There are two problems:

(1) Do I really want to be creating a new data type in global scope and then a new method for this type when I am intending a one-off calculation? Do I want to add minimum 4 lines of code (3 for the struct, parameter, end) when I just want a one-off calculation?
(2) Structs are a pain to use interactively: if I need to tweak my struct definition, I must either (A) restart my session (B) write a new struct with a new name (“my_struct2”).

What would be nice is something like an anonymous functor. The closest you can get to this is defining methods for namedtuples:

(nt::NamedTuple{(:x, :y)})(z::Float64) = sum(z .* nt.y) + nt.x
my_funct = (x = 1, y = [1, 2])
my_funct(10.0) 
# evaluates to 31.0

This works because {NamedTuple:Struct::anonymous function:function}. But it involves type piracy – do I really want to write a method in global scope for all NamedTuples with a particular type signature?

This problem is closely related to discussions of piping syntax. People have suggested creating syntax something like:

f(a, b) = a^b
3 |> f(2, _) 
# returns 8

In the above example, the f(2, _) expression can be thought of as an anonymous closure or an anonymous functor. It should be syntactically feasible to express the idea of a closure as succinctly as above.

13

I’m not sure what the problem that you’re discussing is. You can already do things like (f ∘ g ∘ h)(x).

Is it just that you want to reverse the “flow of information”, that is you want the input to “flow” from left to right, instead of from right to left? In any case, I’m not a fan of syntax sugar, I’d rather see Julia lose some of it than gain more.

14

It’s interesting to note that if bindings are not reassigned, it’s the same effect as and implemented as if the instance were captured. But you can’t always capture instances because closures like do blocks and comprehensions are expected to reassign local variables like any other local scopes; for example, a foreach-do block should do the same thing as a for-loop. Here’s a link to the middle of a thread for a few comments addressing this. Note that as explained prior, the examples’ reassigned variables are Core.Boxed, and the lack of type declarations means retrieving the variable propagates type instability. Annotating the variable, even generically like t::T = 0, isolates the type instability to the Core.Box, though there’s nothing that can be done about the extra heap allocation because a variable that outlives its scope cannot be stored in a stack.

Right, closures handle their underlying functors for you so you don’t have to edit multiple places in code, hence “often a lot more convenient to write”.

That wouldn’t work, an unnamed composite cannot distinguish 2 separate closures that capture the same variables. Closures do make names for the underlying structs, kind of gensymed but not sure how.

1 Like