Performant Recursive Anonymous Functions

[uniment]

What’s the preferred way to write recursive anonymous functions?

Consider a demo function:

julia> function fib(n)
           n ≤ 1 && return n
           return fib(n-1) + fib(n-2)
       end
fib (generic function with 1 method)

julia> @btime fib(10)
  206.726 ns (0 allocations: 0 bytes)
55

Suppose I want to make an “anonymous” version of this function. Of course, the function needs a name to call itself, but maybe I can give it a name in a local scope:

julia> fibby=let
           function fib(n)
               n ≤ 1 && return n
               return fib(n-1) + fib(n-2)
           end
       end
(::var"#fib#1") (generic function with 1 method)

julia> @btime $fibby(10)
  3.725 μs (0 allocations: 0 bytes)
55

It works, but it its performance takes a big hit—running over 10x slower.

It goes type-unstable.

julia> @code_warntype fibby(10)
MethodInstance for (::var"#fib#1")(::Int64)
  from (::var"#fib#1")(n) in Main at REPL[4]:2
Arguments
  #self#::var"#fib#1"
  n::Int64
Locals
  fib@_3::Union{}
  fib@_4::Union{}
Body::Any
1 ─ %1  = (n ≤ 1)::Bool
└──       goto #3 if not %1
2 ─       return n
3 ─ %4  = Core.getfield(#self#, :fib)::Core.Box
│   %5  = Core.isdefined(%4, :contents)::Bool
└──       goto #5 if not %5
4 ─       goto #6
5 ─       Core.NewvarNode(:(fib@_3))
└──       fib@_3
6 ┄ %10 = Core.getfield(%4, :contents)::Any
│   %11 = (n - 1)::Int64
│   %12 = (%10)(%11)::Any
│   %13 = Core.getfield(#self#, :fib)::Core.Box
│   %14 = Core.isdefined(%13, :contents)::Bool
└──       goto #8 if not %14
7 ─       goto #9
8 ─       Core.NewvarNode(:(fib@_4))
└──       fib@_4
9 ┄ %19 = Core.getfield(%13, :contents)::Any
│   %20 = (n - 2)::Int64
│   %21 = (%19)(%20)::Any
│   %22 = (%12 + %21)::Any
└──       return %22

What if we put it inside another function?

julia> function fibby2(n)
           function fib(n)
               n ≤ 1 && return n
               return fib(n-1) + fib(n-2)
           end
           return fib(n)
       end
fibby2 (generic function with 1 method)

julia> @btime fibby2(10)
  3.688 μs (2 allocations: 32 bytes)
55

Same thing.

Is this supposed to happen? What am I missing?

[Sukera]

@uniment

It works, but it its performance takes a big hit—running over 10x slower.

This is because fib needs a reference to itself in a local scope, thereby creating an anonymous function that captures a reference to itself for those subsequent inner calls. It’s essentially equivalent to this:

let
struct var"#fib#1"
   fib::Box
end

function (f::var"#fib#1")(n)
    n <= 1 && return n
    # these are type unstable due to `f.fib` being an untyped `Box`, which also needs to be checked for being defined or not
    return f.fib.contents(n-1) && f.fib.contents(n-2)
end


return var"#fib#1"(<<< this receives itself as its sole argument >>>)
end

Yes, that’s a circular reference at the end. See:

julia> fibby=let
                  function fib(n)
                      n ≤ 1 && return n
                      return fib(n-1) + fib(n-2)
                  end
              end
(::var"#fib#4") (generic function with 1 method)

julia> fibby.fib
Core.Box(var"#fib#4"(Core.Box(#= circular reference @-2 =#)))

The same happens for your fibby2 example. The canonical issue to follow for this and similar performance issues with anonymous functions is https://github.com/JuliaLang/julia/issues/15276

The reason for the Box in the first place is that this transform happens during lowering, way before any type information (or function information, for that matter) exists. Fixing this needs (at least) a rework of the lowering stage of the compiler.

[aplavin]

You can also pass the function as an argument to itself:

julia> fibby=let
                  fib = function(fib, n)
                      n ≤ 1 && return n
                      return fib(fib, n-1) + fib(fib, n-2)
                  end
                  n -> fib(fib, n)
              end
#8 (generic function with 1 method)

julia> @btime fibby(10)
  238.660 ns (0 allocations: 0 bytes)
55

[uniment]

@Sukera

The canonical issue to follow for this and similar performance issues with anonymous functions is https://github.com/JuliaLang/julia/issues/15276

Thanks for the response. I’ve bumped into the fact that lambdas will box the variables they capture from their outside scope before. I guess what surprised me, here, is that I didn’t expect a box to form because I used the syntax which normally avoids a box.

Namely, we have four ways to define a function (named/anonymous and block/inline):

function f(args...) #=...=# end    # syntax A
f(args...) = (#=...=#)             # syntax B
f = function(args...) #=...=# end  # syntax C
f = (args...) -> (#=...=#)         # syntax D

and, at the global scope, a box can be avoided by either a) using A or B, which disallow f from becoming anything other than a Function, or b) pairing C or D with the const keyword.

Example

julia> sad_fib = function(n)
           n ≤ 1 && return n
           return sad_fib(n-1) + sad_fib(n-2)
       end
#42 (generic function with 1 method)

julia> @btime sad_fib(10)
  6.260 μs (0 allocations: 0 bytes)
55

julia> const happy_fib = function(n)
           n ≤ 1 && return n
           return happy_fib(n-1) + happy_fib(n-2)
       end
#44 (generic function with 1 method)

julia> @btime happy_fib(10)
  364.398 ns (0 allocations: 0 bytes)
55

(I ran this in 1.8.0, which is slower than 1.8.3 but faster than 1.9.0.)

Inside the local scope, we don’t have the ability to declare a local variable as const (why?), but we can address type instabilities by declaring the type of the captured variable. This doesn’t always avoid boxing, as it’s possible the value of a captured variable could still change during the function’s lifetime, but if the boxing was due to type instability then it can help.

As I can’t declare the type of a variable which will be set to a lambda before I create the lambda, and I can’t use the const keyword, my instinct went to using syntax A or B.

What is the best way to address this? Is it simply to deprecate syntax A and B inside non-global scopes?

@aplavin

You can also pass the function as an argument to itself:

Ha, I like it!

We can also get *most* of the way there with this.

julia> fibby = let
           function fib(n, fib=fib)
               n ≤ 1 && return n
               return fib(n-1, fib) + fib(n-2, fib)
           end
       end
(::var"#fib#54") (generic function with 2 methods)

julia> @btime $fibby(10)
  489.583 ns (1 allocation: 16 bytes)
55

But I think this is the answer (taking inspiration from this):

julia> fibby_happy = let fib=function(n)
           n ≤ 1 && return n
           return fib(n-1) + fib(n-2)
           end
           fib
       end
#55 (generic function with 1 method)

julia> @btime $fibby_happy(10)
  266.502 ns (0 allocations: 0 bytes)
55

It’s just funny that this is so different:

julia> fibby_sad = let
           local fib=function(n)
               n ≤ 1 && return n
               return fib(n-1) + fib(n-2)
           end
       end
#58 (generic function with 1 method)

julia> @btime $fibby_sad(10)
  6.280 μs (0 allocations: 0 bytes)
55

[sgaure]

There is also the undocumented self-reference var"#self#", but since it’s undocumented you should not rely on it:

fibby = n -> begin
           n ≤ 1 && return n
           return var"#self#"(n-1) + var"#self#"(n-2)
        end

[uniment]

julia> fibby = function(n)
           n ≤ 1 && return n
           return var"#self#"(n-1) + var"#self#"(n-2)
       end
#60 (generic function with 1 method)

julia> @btime $fibby(10)
  336.279 ns (0 allocations: 0 bytes)
55

[uniment]

Oops. This was, in fact, not the solution. It only worked because I previously polluted my global namespace with fib.

[uniment]

Local scope is such a weird place to be

We can define named functions with type-specialized methods,

julia> f = let
           function f(x::Int) x^2 end
           function f(x::Float64) x/2 end
       end
f (generic function with 2 methods)

julia> f(2)
4

julia> f(2.0)
1.0

and then throw them all away

julia> f = let
           function f(x::Int) x^2 end
           function f(x::Float64) x/2 end
           f = 0
       end
0

julia> f
0

It’s so bizarre it feels like this ought to throw an error, the way it would at global scope.

[uniment]

Okay, this is the answer.

julia> fibby = let fib = function fib(n) 
               n≤1 ? n : fib(n-1)+fib(n-2) 
           end 
           fib 
       end
fib (generic function with 1 method)

julia> @btime $fibby(10)
  358.376 ns (0 allocations: 0 bytes)
55

[uniment]

Dang it, that’s not the answer either, because now it’s polluting the global namespace again. Boo.

Wrapping the entire thing in another let block fixes that, but then the type instability returns and we’re back at square one.

[cstjean]

Would using a gensym for the function name help solve the pollution problem for you?

@uniment

Oops. This was, in fact, not the solution. It only worked because I previously polluted my global namespace with fib.

For me it works as well as the global function solution?

julia> fibby = n -> begin
                  n ≤ 1 && return n
                  return var"#self#"(n-1) + var"#self#"(n-2)
               end
#1 (generic function with 1 method)

julia> @btime $fibby(10)
  317.532 ns (0 allocations: 0 bytes)
55

julia> function fib(n)
           n ≤ 1 && return n
           return fib(n-1) + fib(n-2)
       end
fib (generic function with 1 method)

julia> @btime fib(10)
  317.498 ns (0 allocations: 0 bytes)
55

and conceptually, it makes sense, self is the passed function object.

[sgaure]

@uniment

We can define named functions with type-specialized methods,

Here’s a more direct way to define unnamed functions with specialization:

julia> f = n::Integer -> n/2
julia> (::typeof(f))(x::Real) = x^2
julia> f(2)
1.0
julia> f(2.0)
4.0

[Syx_Pek]

One solution is to use a Y-combinator.

julia> fibby = let Y = f -> (x -> x(x))(y -> f((t...) -> y(y)(t...))), 
                   fib = f -> (n -> n ≤ 1 ? n : f(n-1) + f(n-2))
                   
                   Y(fib)
               end;

julia> @btime ($fibby)(10)
  437.399 ns (0 allocations: 0 bytes)
55

[uniment]

What is this black magic!

Unfortunately it’s less performant…

julia> fibby = n -> begin
                         n ≤ 1 && return n
                         return var"#self#"(n-1) + var"#self#"(n-2)
                      end
#15 (generic function with 1 method)

julia> @btime $fibby(10)
  351.442 ns (0 allocations: 0 bytes)
55

julia> fibby = let Y = f -> (x -> x(x))(y -> f((t...) -> y(y)(t...))),
                          fib = f -> (n -> n ≤ 1 ? n : f(n-1) + f(n-2))

                          Y(fib)
                      end;

julia> @btime $fibby(10)
  505.699 ns (0 allocations: 0 bytes)
55

@cstjean

Would using a gensym for the function name help solve the pollution problem for you?

Yes it would. Although I think I mis-named this thread; I actually would prefer the function to have a name, just a local name. Sorry for being misleading.

@cstjean

For me it works as well as the global function solution?

Sorry I was responding to myself; I thought I had found a solution that didn’t involve using gensym or var"#self#". As @sgaure mentioned, I probably don’t want to use var"#self#" anyway.

[uniment]

And this just causes Julia to crash:

julia> fibby = let
           function f end
           f::typeof(f) = function(n)
               n ≤ 1 ? n : f(n-1)+f(n-2)
           end
       end

as does this:

julia> fibby = let
           function f end
           f::typeof(f) = f
           f(n) = n ≤ 1 ? n : f(n-1)+f(n-2)
           f
       end

@sgaure

Here’s a more direct way to define unnamed functions with specialization:

julia> f = n::Integer -> n/2
julia> (::typeof(f))(x::Real) = x^2

What’s the implication?

Unnamed functions can be overloaded too, and therefore there’s nothing special about named functions, and therefore if we can’t have consts inside local scopes then we shouldn’t expect to have named functions behave as they do at global scope w.r.t. type-stability?

And therefore, we should probably just deprecate named functions named function syntax inside local scopes to avoid confusion?

[uniment]

Holy smokes, I think we found a winner. Thanks @sgaure for the inspiration!

julia> fibby = let
           function f end
           (f::typeof(f))(n) = n ≤ 1 ? n : f(n-1)+f(n-2)
           f
       end
(::var"#f#13") (generic function with 1 method)

julia> @btime fibby(10)
  388.500 ns (0 allocations: 0 bytes)
55

[Mason]

Yeah, this is a funny, and very aggravating property of how we handle anonymous functions. One other solution I’ve seen, first shown to me by @tkf is this:

julia> using MacroTools, ExprTools

julia> fibby = let
           function _fib(f, n)
               n ≤ 1 ? n : f(f, n-1) + f(f, n-2)
           end
           fib(n) = _fib(_fib, n)
       end

julia> @btime $fibby($10)
  221.356 ns (0 allocations: 0 bytes)
55

The idea here is to basically just play some sleight of hand to avoid julia thinking we’re doing recursion, and to stop it from doing something dumb.

This can be automated into a macro with a little elbow grease:

using ExprTools: splitdef, combinedef
using MacroTools: postwalk, @capture

"""
This is only worth using for locally scoped recursive functions.
"""
macro fast_recustion(fdef)
    d = splitdef(fdef)
    name = get!(d, :name, gensym(:f))
    fargs = copy(get!(d, :args, []))
    fkwargs = get!(d, :kwargs, [])
    _name = gensym(name)
    __name = gensym(Symbol(:_, name))
    d[:body] = postwalk(d[:body]) do ex
        if @capture(ex, f_(args__))
            if f == name
                return :($__name($__name, $(args...),))
            end
        elseif @capture(ex, f(args__; kwargs__))
            if f == name
                return :($__name($__name, $(args...); $(kwargs...),))
            end
        end
        ex
    end
    d[:name] = _name
    d[:args] = pushfirst!(d[:args], __name)
    quote
        $(combinedef(d))
        $name($(fargs...); $(fkwargs...),) = $_name($_name, $(fargs...); $(fkwargs...),)
    end |> esc
end

so that we can write

julia> fibby = let
           @fast_recustion function fib(n)
               n ≤ 1 ? n : fib(n-1) + fib(n-2)
           end
       end
(::var"#fib#56"{var"##fib#456#55"}) (generic function with 1 method)

julia> @btime $fibby($10)
  218.555 ns (0 allocations: 0 bytes)
55

[uniment]

That’s what @aplavin was suggesting, but a little nicer. What I wasn’t a fan of, is that now you’re compiling two functions:

julia> fibby = let
           function f end
           (f::typeof(f))(n) = n ≤ 1 ? n : f(n-1)+f(n-2)
           f
       end
(::var"#f#18") (generic function with 1 method)

julia> @time fibby(10)
  0.007368 seconds (1.66 k allocations: 87.968 KiB, 99.62% compilation time)
55

julia> fibby = let
                  function _fib(f, n)
                      n ≤ 1 ? n : f(f, n-1) + f(f, n-2)
                  end
                  fib(n) = _fib(_fib, n)
              end
(::var"#fib#20"{var"#_fib#19"}) (generic function with 1 method)

julia> @time fibby(10)
  0.011338 seconds (2.91 k allocations: 150.917 KiB, 99.75% compilation time)
55

[uniment]

For fun, setting f and g to have separate names.

Notice the subtle difference, between calling f (which is an argument to the function), versus calling g (which is not).

julia> fibby = let
           function g end
           (f::typeof(g))(n) = n ≤ 1 ? n : f(n-1)+f(n-2)
           g
       end
(::var"#g#25") (generic function with 1 method)

julia> @btime $fibby(10)
  336.181 ns (0 allocations: 0 bytes)
55

julia>

julia> fibby = let
           function g end
           (f::typeof(g))(n) = n ≤ 1 ? n : g(n-1)+g(n-2)
           g
       end
(::var"#g#26") (generic function with 1 method)

julia> @btime $fibby(10)
  6.160 μs (0 allocations: 0 bytes)
55

[Sukera]

Pretty much all of the examples that work with the requested speed here do the same thing: not having the anonymous function refer to itself in its body, avoiding the Box around itself. Some versions refer to the singleton instance (the implicit first argument) of the function (that’s the case for (f::typeof(g))(n) = n <= 1 ? n : f(n-1)+f(n-2)), some use a second function (that’s the Y-combinator solution), some avoid the Box by passing the function to recurse on explicitly as an argument (that’s the approaches with passing the recursive function explicitly as the first argument).

@uniment

I actually would prefer the function to have a name, just a local name.

If I understand what you want correctly, this is not a thing - all function definitions happen at the containing global scope. There are no “local only” functions. They ALWAYS get a name at the global level, because that’s relevant for world age.

@uniment

Notice the subtle difference, between calling f (which is an argument to the function), versus calling g (which is not).

The difference again is the self boxing, due to g not being known to be anything in particular during lowering, as mentioned in my first reply. It’s the exact same behavior.

I encourage you to check the various versions in here with Meta.@lower, if you’re not afraid of looking at lowered julia code without type information, as the compiler first sees it.