Optimizing a recursive function with metaprogramming in Julia

Question

Following the approach of this answer I am trying to understand what happens exactly and how expressions and generated functions work in Julia within the concept of metaprogramming.

The goal is to optimize a recursive function using expressions and generated functions (for a concrete example you can have a look at the question answered in the link provided above).

Consider the following modified fibonacci function, in which I want to compute the fibonacci series up to n and multiply it by a number p.

The straightforward, recursive implementation would be

function fib(n::Integer, p::Real)
    if n <= 1
        return 1 * p
    else
        return n * fib(n-1, p)
    end
end

As a first step, I could define a function which returns an expression instead of the computed value

function fib_expr(n::Integer, p::Symbol)
    if n <= 1
        return :(1 * $p)
    else
        return :($n * $(fib_expr(n-1, p)))
    end
end

which, e.g. returns something like

julia> ex = fib_expr(3, :myp)
:(3 * (2 * (1myp)))

In this way I get an expression which is fully expanded and depends on the value assigned to the symbol myp. In this way I do not see the recursion anymore, basically I am metaprogramming: I created a function that creates another "function" (in this case we call it expression though). I can now set myp = 0.5 and call eval(ex) to compute the result. However, this is slower than the first approach.

What I can do though, is to generate a parametric function in the following way

@generated function fib_gen{n}(::Type{Val{n}}, p::Real)
    return fib_expr(n, :p)
end

And magically, calling fib_gen(Val{3}, 0.5) gets things done, and is incredibly fast.
So, what is going on?

To my understanding, in the first call to fib_gen(Val{3}, 0.5), the parametric function fib_gen{Val{3}}(...) gets compiled and its content is the fully expanded expression obtained through fib_expr(3, :p), i.e. 3*2*1*p with p substituted with the input value. The reason why it is so fast then, is because fib_gen is basically just a series of multiplications, whereas the original fib has to allocate on the stack every single recursive call making it slower, am I correct?

To give some numbers, here is my short benchmark using BenchmarkTools.

julia> @benchmark fib(10, 0.5)
...
mean time: 26.373 ns
...

julia> p = 0.5
0.5

julia> @benchmark eval(fib_expr(10, :p))
...
mean time: 177.906 μs
...

julia> @benchmark fib_gen(Val{10}, 0.5)
...
mean time: 2.046 ns
...

I have many questions:

Why the second case is so slow?
What exactly is and means ::Type{Val{n}}? (I copied that from the answer linked above)
Because of the JIT compiler, sometimes I am lost in what happens at compile-time and at run-time, as it is the case here...

Furthermore, I tried to combine fib_expr and fib_gen in a single function according to

@generated function fib_tot{n}(::Type{Val{n}}, p::Real)
    if n <= 1
        return :(1 * p)
    else
        return :(n * fib_tot(Val{n-1}, p))
    end
end

which however is slow

julia> @benchmark fib_tot(Val{10}, 0.5)
...
mean time: 4.601 μs
...

What am I doing wrong here? Is it even possible to combine fib_expr and fib_gen in a single function?

I realize this is more a monograph rather than a question, however, even though I read the metaprogramming section few times, I am having a hard time to grasp everything, in particular with an applied example such as this one.

I like the question, and not only because I triggered it -- very well written. However, using Fibonacci for generated functions is probably not the best example in reality: in your other question, the values for which the methods were generated were bounded beforehand within a small range (by convention), whereas here, you provide the possibility to compile a new one for each single, unbounded value. But it's good enough for understanding the concept, I guess :) — phipsgabler
I agree it is not exactly the same in practice, however for understanding the concept I believe it considers all points touched in my other question. On top of that, in principle, from what I understood so far, fib_gen is actually going to be always faster than fib as long as generating the expression(s) does not provoke a StackOverflowError. — stefabat
If you call it a lot of times, yes. But the first call will be much slower than one call of just fib. It depends on your use case whether it pays off, and for Fibonacci numers, it probably won't often. That's what I meant to say. — phipsgabler
You can metaprogram in the language, or from the outside. If you do it inside, you have to understand what the language/compiler offers and its strengths and limitations, often many. If you do from outside, a) there are no limitations due the language or compiler, b) the concepts are often cleaner because they have to be general. See my Software Engineering answer on metaprogramming: softwareengineering.stackexchange.com/a/257441/12135 — Ira Baxter

phipsgabler phipsgabler · Accepted Answer · 2018-07-06T13:17:16

A monograph in response:

Metaprogramming basics

It will be easier to start with "normal" macros first. I'll relax the definition you used a bit:

function fib_expr(n::Integer, p)
    if n <= 1
        return :(1 * $p)
    else
        return :($n * $(fib_expr(n-1, p)))
    end
end

That allows to pass in more than just symbols for p, like integer literals or whole expressions. Given this, we can define a macro for the same functionality:

macro fib_macro(n::Integer, p)
    fib_expr(n, p)
end

Now, if @fib_macro 45 1 is used anywhere in the code, at compile time it will first be replaced by a long nested expression

:(45 * (44 * ... * (1 * 1)) ... )

and then compiled normally -- to a constant.

That's all there is to macros, really. Replacing syntax during compile time; and by recursion, this can be an arbitrarily long alteration between compiling, and evaluating functions on expressions. And for things that are essentially constant, but tedious to write otherwise, it is very useful: a bood example example is Base.Math.@evalpoly.

Evaluation at runtime?

But it has the problem that you cannot inspect values which are only known at runtime: you can't implement fib(n) = @fib_macro n 1, since at compile time, n is a symbol representing the parameter, and not a number you can dispatch on.

The next best solution to this would be to use

fib_eval(n::Integer) = eval(fib_expr(n, 1))

which works, but will repeat the compilation process every time it is called -- and that is much more overhead than the original function, since now at runtime, we perform the whole recursion on the expression tree and then call the compiler on the result. Not good.

Method dispatch & compilation

So we need a way to intermingle runtime and compile time. Enter @generated functions. These will at runtime dispatch on a type, and then work like a macro defining the function body.

First about type dispatch. If we have

f(x) = x + 1

and have a function call f(1), about the following will happen:

The type of the argument is determined (Int)
The method table of the function is consulted to find the best matching method
The method body is compiled for the specific Int argument type, if that hasn't been done before
The compiled method is evaluated on the concrete argument

If we then enter f(1.0), the same will happen again, with a new, different specialized method being compiled for Float64, based on the same function body.

Value types & singleton types

Now, Julia has the peculiar feature that you can use numbers as types. That means that the dispatch process outlined above will also work on the following function:

g(::Type{Val{N}}) where N = N + 1

That's a bit tricky. Remember that types are themselves values in Julia: Int isa Type.

Here, Val{N} is for every N a so-called singleton type having exactly one instance, namely Val{N}() -- just like Int is a type having many instances 0, -1, 1, -2, ....

Type{T} is also a singleton type, having as its single instance the type T. Int is a Type{Int}, and Val{3} is a Type{Val{3}} -- in fact, both are the only values of their type.

So, for each N, there is a type Val{N}, being the single instance of Type{Val{N}}. Thus, g will be dispatched and compiled for each single N. This is how we can dispatch on numbers as types. This already allows for optimization:

julia> @code_llvm g(Val{1})

define i64 @julia_g_61158(i8**) #0 !dbg !5 {
top:
  ret i64 2
}

julia> @code_llvm f(1)

define i64 @julia_f_61076(i64) #0 !dbg !5 {
top:
  %1 = shl i64 %0, 2
  %2 = or i64 %1, 3
  %3 = mul i64 %2, %0
  %4 = add i64 %3, 2
  ret i64 %4
}

But remember that it requires compilation for each new N at the first call.

(And fkt(::T) is just short for fkt(x::T) if you don't use x in the body.)

Integrating generating functions and value types

Finally to generated functions. They work as a slight modification of the above dispatch pattern:

The type of the argument is determined (Int)
The method table of the function is consulted to find the best matching method
The method body is treated as a macro and called with the Int argument type as a parameter, if that hasn't been done before. The resulting expression is compiled into a method.
The compiled method is evaluated on the concrete argument

This pattern allows to change the implementation for each type which the function is dispatched on.

For our concrete setting, we want to dispatch on the Val types representing the arguments of the Fibonacci sequence:

@generated function fib_gen{n}(::Type{Val{n}}, p::Real)
    return fib_expr(n, :p)
end

You now see that your explanation was exactly right:

in the first call to fib_gen(Val{3}, 0.5), the parametric function fib_gen{Val{3}}(...) gets compiled and its content is the fully expanded expression obtained through fib_expr(3, :p), i.e. 3*2*1*p with p substituted with the input value.

I hope that the whole story has also answered all three of your listed questions:

The implementation using eval replicates the recursion every time, plus the overhead of compilation
Val is a trick to lift numbers to types, and Type{T} the singleton type containing only T -- but I hope the examples were helpful enough
Compile time is not before execution, because of JIT -- it is every time a method gets compiled first time, because it get's called.