Make EvalFunc hold world age and deprecate GeneralizedGenerated

[ChrisRackauckas]

@c42f's PR JuliaLang/julia#35844 gives an invoke_in_world that we could use to replace how we've currently been using GeneralizedGenerated in a way that is safer for general Julia code. Currently our build_function has two routes: one that uses eval and generates an EvalFunc, and the other which builds the function via GeneralizedGenerated. The eval version had to be created because the type from GeneralizedGenerated can break. If you have a million equations, you get a gigantic type and that just seems to do some outrageous unholy things to the compiler.

EvalFunc requires one to invoke in the latest, so the machinery is defined in DiffEqBase so that we have

https://github.com/SciML/DiffEqBase.jl/blob/master/src/solve.jl#L5-L8

and then https://github.com/SciML/DiffEqBase.jl/blob/master/src/solve.jl#L30 which basically gives us a hook to get around world age issues when GeneralizedGenerated is not used. This works but has all of the downsides of invokelatest.

However, this new invoke is definitely a good way forward. For example, when we generate the functions from the symbolic code, i.e. https://github.com/SciML/ModelingToolkit.jl/blob/master/src/systems/diffeqs/abstractodesystem.jl#L122-L152, instead of building just an EvalFunc, we can make a WorldFunc that capturers the WorldAge as type information. Then we can have a call on this object perform invoke_in_world at the right world age. Hopefully, storing the world age as type information will make this all inferred and just as fast as a standard Julia function? (that's a question for @c42f).

With this in place, we can remove GeneralizedGenerated usage because then the standard function would be safe due to invoking in the proper world. The only issue that would come up is then that this is not precompile safe because the generated code is made for a specific world. However, any function that is made to precompile is a function evaluated in a module scope, which means it's safe to just use eval and not invoke in the world, and so we just need a flag to turn off the WorldFunc wrapping. That should make the whole symbolic system generate functions that are both efficient and don't have world age issues, and not have the scaling issues of GG.

[c42f]

Somehow I missed this question, sorry!

we can make a WorldFunc that capturers the WorldAge as type information. Then we can have a call on this object perform invoke_in_world at the right world age. Hopefully, storing the world age as type information will make this all inferred and just as fast as a standard Julia function? (that's a question for @c42f).

This was something I was originally hoping for. But IIUC it turns out not to work as a replacement for GeneralizedGenerated; you'd still need an invokelatest boundary somewhere because the compiler can't see a newer world while compiling for an older one.

We did discuss this four months ago on the compiler channel. The conclusion was that a compiler-friendly replacement for GeneralizedGenerated would be callable methods. Here's what I took away from the discussion at the time:

So to summarize this thread, a possible solution for the "turn this symbolically computed AST into a callback" use case common to ModellingToolkit, Soss, etc seems to be a way to opt out of dispatch (and hence the world mechanism) when creating the callback. Either by having callable Methods or some better/alternative abstraction. Maybe YAKCs if they apply to this use case.

For your immediate use case, can you just do the UUID workaround: cache the AST in a global table with a new UUID as key which is put into a type as a proxy for the AST. Then pull the source from that cache within a generated function?

[ChrisRackauckas]

Hmm, I'm getting overhead when I try this:

const function_cache = Dict{UInt64,Expr}()
struct GeneratedFunction{uuid,argnames} end
(f::GeneratedFunction{uuid})(args...) where uuid = generated_callfunc(f,args...)

@generated function generated_callfunc(f::GeneratedFunction{uuid,argnames},__args...) where {uuid,argnames}
    setup = (:($(argnames[i]) = @inbounds __args[$i]) for i in 1:length(argnames))
    quote
        $(setup...)
        $(function_cache[uuid])
    end
end

function f(_du,_u,_p,_t)
    @inbounds _du[1] = _u[1]
    @inbounds _du[2] = _u[2]
    nothing
end
ex = :(function f(_du,_u,_p,_t)
    @inbounds _du[1] = _u[1]
    @inbounds _du[2] = _u[2]
    nothing
end)

function_cache[hash(ex.args[2])] = ex.args[2]
args = (ex.args[1].args[2:end]...,)

myfunc = GeneratedFunction{hash(ex.args[2]),args}()

du = rand(2)
u = rand(2)
p = nothing
t = nothing
myfunc(du,u,p,t)
du

using BenchmarkTools
@btime myfunc(du,u,p,t) # 195.563 ns (1 allocation: 32 bytes)
@btime f(du,u,p,t) # 18.035 ns (0 allocations: 0 bytes)

[ChrisRackauckas]

The generated code looks right:

uuid = hash(ex.args[2])
argnames = args
setup = (:($(argnames[i]) = @inbounds __args[$i]) for i in 1:length(argnames))

quote
    $(setup...)
    $(function_cache[uuid])
end

quote
    #= D:\OneDrive\Computer\Desktop\test.jl:75 =#
    _du = #= D:\OneDrive\Computer\Desktop\test.jl:72 =# @inbounds(__args[1])
    _u = #= D:\OneDrive\Computer\Desktop\test.jl:72 =# @inbounds(__args[2])
    _p = #= D:\OneDrive\Computer\Desktop\test.jl:72 =# @inbounds(__args[3])
    _t = #= D:\OneDrive\Computer\Desktop\test.jl:72 =# @inbounds(__args[4])
    #= D:\OneDrive\Computer\Desktop\test.jl:76 =#
    begin
        #= D:\OneDrive\Computer\Desktop\test.jl:59 =#
        #= D:\OneDrive\Computer\Desktop\test.jl:60 =#
        #= D:\OneDrive\Computer\Desktop\test.jl:60 =# @inbounds _du[1] = _u[1]
        #= D:\OneDrive\Computer\Desktop\test.jl:61 =#
        #= D:\OneDrive\Computer\Desktop\test.jl:61 =# @inbounds _du[2] = _u[2]
        #= D:\OneDrive\Computer\Desktop\test.jl:62 =#
        nothing
    end
end

[YingboMa]

I think the allocation comes from getfields

julia> @code_typed debuginfo=:none myfunc(du,u,p,t)
CodeInfo(
1 ─ %1 = (getfield)(args, 1)::Array{Float64,1}
│   %2 = (getfield)(args, 2)::Array{Float64,1}
│   %3 = Base.arrayref(false, %2, 1)::Float64
│        Base.arrayset(false, %1, %3, 1)::Array{Float64,1}
│   %5 = Base.arrayref(false, %2, 2)::Float64
│        Base.arrayset(false, %1, %5, 2)::Array{Float64,1}
│   %7 = Main.nothing::Nothing
└──      return %7
) => Nothing

[ChrisRackauckas]

@YingboMa figured out that it was just specialization:

@btime $myfunc($du,$u,$p,$t) # 2.099 ns (0 allocations: 0 bytes)
@btime f($du,$u,$p,$t) # 2.099 ns (0 allocations: 0 bytes)

Just like anonymous functions. So this is being created in https://github.com/SciML/RuntimeGeneratedFunctions.jl and MTK will make use of the package.

[c42f]

RuntimeGeneratedFunctions looks about right to me. Couple of thoughts:

• Did you need the cache to be content-addressed? If not, you could just use UUIDs.uuid4() as a key and skip the hashing.
• If you need the function cache to survive precompilation, you can do this by putting it into the module which creates the RuntimeGeneratedFunction rather than putting it into RuntimeGeneratedFunctions itself. (That is, you have one function cache per module.) You can use a hidden variable name like var"#_function_cache if you don't want users to see it in their module.

[ChrisRackauckas]

If you need the function cache to survive precompilation, you can do this by putting it into the module which creates the RuntimeGeneratedFunction rather than putting it into RuntimeGeneratedFunctions itself. (That is, you have one function cache per module.) You can use a hidden variable name like var"#_function_cache if you don't want users to see it in their module.

Oh interesting! Yes, that would be extremely useful, since I was just about to implement this in a way where there's a separate hook to use eval to then document a precompile-safe way: if this is precompile safe then it makes everything a lot easier.

[c42f]

Depends what you mean by precompile-safe.

But yes, the function cache should survive precompilation. A tricky thing is to make sure that RuntimeGeneratedFunctions itself can look up the cache in the other modules. The modules which use it might need some code in __init__() to make this work.

I did some similar things when fixing up Unitful's @u_str macro so that it was able to see the units defined in unit extension modules.

[ChrisRackauckas]

If you could lend a hand for that sometime it would be helpful. It doesn't need to be "soon".

[c42f]

Happy to. So by "precompile-safe", I assume you mean that RuntimeGeneratedFunctions which are generated by package top-level code (or during macro expansion) will work correctly when the module is precompiled and later loaded?

[ChrisRackauckas]

yes thanks!