Dynamically Create a Function - Initial Idea with eval failed due to World Age Issue

Hey all,

I’m struggling to implement a certain problem.
My current situation is the following: Some parts of my project consist of using Tullio functions. Since my arrays can be of arbitrary dimensions and I want to apply some user specific operations on them, I was creating functions with meta programming to compose the different Tullio expressions.

This minimal working example somehow shows my use case:

# function to create a complex expression
function N_dim(dim)
    if dim == 1
        return Meta.parse("arr[1, 1] .+ arr[1, 2]")
    else
        return Meta.parse("arr[1] .+ arr[2]")
    end 
end

# convenient wrapper to create a function from that expression 
function make_f(dim)
    f = @eval arr -> $(N_dim(dim))
    return f
end

# a function from which our f should be called
function foo()
    f = make_f(1)
    @show f([1 2; 3 4])    # fails here
    return 0
end

foo()

As known, this case runs into the world age issue:

ERROR: MethodError: no method matching (::var"#21#22")(::Int64)
The applicable method may be too new: running in world age 27854, while current world is 27855.

I somehow grasped the issue of that and already read a lot of regarded posts. However, in my case I can’t really use Base.invokelatest because I want to pass that function into Zygote.jl. Furthermore this function is performance critical and should be as fast as possible (hence the reason for using Tullio).

I would be really happy if someone could help me in solving that issue. Generating functions for different inputs at compile time is not really feasible since my real make_f has several input parameters.

Thanks,

Felix

Have you seen this thread:

I am happy with

GeneralizedGenerated

and all my tests I have shown that it results in best performance for generated functions.

2 Likes

Thanks for the link, looks promising.
I must have somehow missed it :roll_eyes:

1 Like

I guess it works:

using GeneralizedGenerated

function N_dim(dim)
    if dim == 1
        return Meta.parse("arr -> arr[1, 1] .+ arr[1, 2]")
    else
        return Meta.parse("arr -> arr[1] .+ arr[2]")
    end 
end


function make_f(dim)
    f = mk_function(N_dim(dim))
    return f
end


function foo()
    @time f = make_f(1)
    x = randn((1000, 1000))
    @time f(x)
    @time f(x)
    @time f(x)

    @time f = make_f(2)
    @time f(x)
    @time f(x)
    return f

end

g = foo()

Output:

 2.923102 seconds (4.99 M allocations: 278.121 MiB, 4.95% gc time)
  0.213862 seconds (237.62 k allocations: 13.494 MiB)
  0.000002 seconds (1 allocation: 16 bytes)
  0.017559 seconds (14.34 k allocations: 779.257 KiB)
  0.080903 seconds (62.61 k allocations: 3.381 MiB)
  0.000002 seconds (1 allocation: 16 bytes)

So the first creation takes a while and I’m also surprised how much it allocates.

My final question would be: is that really the solution to that?

Thanks,

Felix

1 Like

Don’t know, got no different solution in my old thread.

1 Like

Turns out that the naive approach doesn’t work with GeneralizedGenerated.

See on GitHub.

You need Base.invokelatest.

julia> function N_dim(dim)
           if dim == 1
               return Meta.parse("arr[1, 1] .+ arr[1, 2]")
           else
               return Meta.parse("arr[1] .+ arr[2]")
           end
       end
N_dim (generic function with 1 method)

julia> function make_f(dim)
           f = @eval arr -> $(N_dim(dim))
           return x->Base.invokelatest(f, [1 2; 3 4]) # CHANGED HERE
       end
make_f (generic function with 1 method)

julia> function foo()
           f = make_f(1)
           @show f([1 2; 3 4])    # no longer fails here
           return 0
       end
foo (generic function with 1 method)
https://arxiv.org/abs/2010.07516
julia> foo()
f([1 2; 3 4]) = 3
0

If you want an explanation for why this is necessary, see this preprint on arXiv:

Did you try GitHub - SciML/RuntimeGeneratedFunctions.jl: Functions generated at runtime without world-age issues or overhead ? GeneralizedGenerated.jl builds massive types. This is usually a bit leaner and causes less related segfaults.

2 Likes

@ChrisRackauckas Tullio generates closures…

However there are still some obstacles to using GG for Tullio.

Hence, unfortunately, I don’t think there is an existing approach for Tullio to avoid world age problem by now…

Advice:

  1. Slightly modifying the code generator to avoid explicit annotations and use GG.
  2. Try eliminating closures to use RuntimeGeneratedFunctions.
    This is possible, because after a short investigation on Tullio generated code, I find that nested functions never escape, and all closures are used locally.
1 Like

Oh… yeah that’s going to give RuntimeGeneratedFunctions.jl issues. Why are the closures needed though? What’s it enclosing? If it’s enclosing something, it could runtime eval the inner function and enclose data on the runtime eval’d version.

You can use multiple dispatch for this example, no need for metaprogramming:

f(::Val{1}) = arr -> arr[1, 1] .+ arr[1, 2]
f(::Any) = arr -> arr[1] .+ arr[2]

f(Val(dim)) # returns the 1st definition if dim=1, otherwise the 2nd one

Alternatively, just a regular higher-order function which returns the functions above, without needing @eval.

Or do I miss something?

2 Likes

My functions are unfortunately much more general than my simple example.

function N_dim(dim, offset)
    if dim == 1
        return Meta.parse("arr -> @tullio res = arr[i] + arr[i+offset]")
    else
        return Meta.parse("arr -> @tullio res = arr[i, j] + arr[i+offset,j ] + arr[i, j] + arr[i,j + offset]")
    end 
end

More information for the generated code is also in this issue.

I got a case with a crash of Julia with GeneralizedGenerated.jl, so I’m going to try RuntimeGeneratedFunctions.jl in my little project to see, how this is working. Thanks for the tip.

Zygote doesn’t work with Base.invokelatest. I maybe ask that in a Zygote issue (or someone here know the answer)?

julia> using Zygote

julia> function f()
           g = @eval x -> $(:(x+1))
           Zygote.gradient(Base.invokelatest(g, 10))
       end
f (generic function with 1 method)

julia> f()
ERROR: MethodError: objects of type Int64 are not callable
Stacktrace:
 [1] macro expansion at /home/fxw/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0 [inlined]
 [2] _pullback(::Zygote.Context, ::Int64) at /home/fxw/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:12
 [3] _pullback(::Int64) at /home/fxw/.julia/packages/Zygote/c0awc/src/compiler/interface.jl:38
 [4] pullback(::Int64) at /home/fxw/.julia/packages/Zygote/c0awc/src/compiler/interface.jl:44
 [5] gradient(::Int64) at /home/fxw/.julia/packages/Zygote/c0awc/src/compiler/interface.jl:53
 [6] f() at ./REPL[6]:3
 [7] top-level scope at REPL[7]:1
 [8] run_repl(::REPL.AbstractREPL, ::Any) at /build/julia/src/julia-1.5.2/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:288

Furthermore I’m afraid of the performance drop due to it.

Thanks for your support (here and Github)!

  1. Slightly modifying the code generator to avoid explicit annotations and use GG.
    You mean the Tullio code generator?
  1. Try eliminating closures to use RuntimeGeneratedFunctions.
    This is possible, because after a short investigation on Tullio generated code, I find that nested functions never escape, and all closures are used locally.

That’s also regarding Tullio, isn’t it?

Try Base.invokelatest(Zygote.gradient, g, 10) instead. Zygote.gradient expects a function, you are passing it just the result of Base.invokelatest(g, 10) though.

4 Likes

Thanks, that’s obviously true :sweat_smile:

Here MWE:

using Tullio
using Zygote
using BenchmarkTools

function f() 
    g = @eval arr -> $(:(@tullio r = (arr[i,j] + arr[i+1, j+1] + arr[i-1, j] + arr[i, j-1])^2))
    m = arr -> @tullio r = (arr[i,j] + arr[i+1, j+1] + arr[i-1, j] + arr[i, j-1])^2
    x = randn((1000, 1000))
    @btime Base.invokelatest(Zygote.gradient, $g, $x)

    @btime Zygote.gradient($m, $x)
end


f()

Output:

  2.701 ms (164 allocations: 7.64 MiB)
  2.648 ms (149 allocations: 7.64 MiB)

So it seems to work.

I’ll try that in my project and report back whether performance issues or errors come up

EDIT:
It works in my larger project. But you need to put the Base.invokelatest at the very outer call.
For example this doesn’t work:

using Tullio
using Zygote
using BenchmarkTools

function f() 
    g = @eval arr -> $(:(@tullio r = (arr[i,j] + arr[i+1, j+1] + arr[i-1, j] + arr[i, j-1])^2))
    m = arr -> @tullio r = (arr[i,j] + arr[i+1, j+1] + arr[i-1, j] + arr[i, j-1])^2
    x = randn((512, 512))
    @btime Zygote.gradient(x -> Base.invokelatest($g, x), $x)

    @btime Zygote.gradient($m, $x)
end

f()

Because then Zygote wants to do the derivative of Base.invokelatest

I think it should be possible to add a rule for Base.invokelatest, but not 100% sure how this plays out with Zygote’s internals.

2 Likes

I have the same point here. You might check How to implement n-dimensional functions? · Issue #11 · mcabbott/Tullio.jl · GitHub

some code generated by Tullio
 begin
            #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:962 =#
            let 𝒜𝒸𝓉! = 𝒜𝒸𝓉!
                #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:963 =#
                begin
                    #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:956 =#
                    local function ℳ𝒶𝓀ℯ(arr)
                            $(Expr(:meta, :inline))
                            #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:956 =#
                            #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:957 =#
                            local 𝒶𝓍i = intersect(eachindex(arr), eachindex(arr) .- 1)
                            local 𝓇𝒽𝓈(arr, i) = begin
                                        #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:780 =#
                                        identity((arr[i] - arr[i + 1]) ^ 2)
                                    end
                            begin
                                #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:789 =#
                                local 𝒯1 = Core.Compiler.return_type(𝓇𝒽𝓈, typeof((arr, first(𝒶𝓍i))))
                                #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:790 =#
                                local 𝒯2 = if Base.isconcretetype(𝒯1)
                                            #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:791 =#
                                            𝒯1
                                        else
                                            #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:793 =#
                                            nothing
                                            #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:794 =#
                                            typeof(𝓇𝒽𝓈(arr, first(𝒶𝓍i)))
                                        end
                            end
                            local 𝒯3 = 𝒯2
                            local 𝒯 = 𝒯3
                            local ℛ = convert(𝒯, zero(𝒯))
                            (Tullio.thread_scalar)(𝒜𝒸𝓉!, (Tullio.storage_type)(arr), ℛ, tuple(arr), tuple(𝒶𝓍i), +, 87382, nothing)
                        end
                end
                #= C:\Users\twshe\.julia\packages\Tullio\Q2RDS\src\macro.jl:964 =#
                ((Tullio.Eval)(ℳ𝒶𝓀ℯ, ∇ℳ𝒶𝓀ℯ))(arr)
            end
        end

You can see Tullio.thread_scalar accepts the generated closure…
I guess everything is possible to get “statically” inlined before execution, but it might be difficult for people to implement.

Suddenly I feel that we might need something like high order macros and local macros.

No time today but perhaps this page of docs is useful, explaining what functions it creates: https://gist.github.com/mcabbott/7494c8cc46a906fe88172dbfc8b92ad2

I’m never too sure of the terminology, they don’t close over values (which they explicitly receive as arguments) but I guess the outer functions “make” close over the kernel function “act!”.