I have a project where I generate a lot of functions at runtime, use RuntimeGeneratedFunctions.jl on them, and store them in a dict for easy access. Then I wrote small wrapper functions to make typing them out easier.
I was benchmarking and it looks like this organization scheme creates a lot of allocations on every function call.
I checked @code_warntype
, and sure enough some of the wrapper functions I made generated type uncertainty. Here's a quick demo:
a(x) = x^2
dict = Dict("a" => a)
@code_warntype dict["a"](1.0)
# no type uncertainty:
#
# Variables
# #self#::Core.Compiler.Const(a, false)
# x::Float64
#
#Body::Float64
#1 ─ %1 = Core.apply_type(Base.Val, 2)::Core.Compiler.Const(Val{2}, false)
#│ %2 = (%1)()::Core.Compiler.Const(Val{2}(), false)
#│ %3 = Base.literal_pow(Main.:^, x, %2)::Float64
#└── return %3
b(x) = dict["a"](x)
@code_warntype b(1.0)
# type uncertainty!
# Variables
# #self#::Core.Compiler.Const(b, false)
# x::Float64
#
#Body::Any
#1 ─ %1 = Base.getindex(Main.dict, "a")::Any
#│ %2 = (%1)(x)::Any
#└── return %2
Is there a way to reduce the type uncertainty of b
, or swap it for another notationally convenient object without type uncertainty?
Do you constantly add key-value pairs dynamically in your use case? If not (i.e. if you can previously assume the amount of generated functions), I think you could just use a NamedTuple instead of a Dict.
(Quick disclaimer: I have no experience with functions generated at runtime.)
This example mask the problem, because type instability in this case originates from the fact that dict
is a global variable. So if you use const dict = Dict("a" => a)
then you'll see no type instabilities.
But real problem is different and can be seen if we are introduce more functions:
a1(x) = x^2
a2(x) = x^3
a3(x) = x^4
a4(x) = x^5
a5(x) = x^6
julia> const dict = Dict("a1" => a1, "a2" => a2, "a3" => a3, "a4" => a4, "a5" => a5)
Dict{String, Function} with 5 entries:
In this case, dict
is a dict of functions and their types can't be predicted at compile time.
julia> b(x) = dict["a1"](x)
b (generic function with 1 method)
julia> @code_warntype b(1.0)
MethodInstance for b(::Float64)
from b(x) in Main at REPL[13]:1
Arguments
#self#::Core.Const(b)
x::Float64
Body::Any
1 ─ %1 = Base.getindex(Main.dict, "a1")::Function
│ %2 = (%1)(x)::Any
└── return %2
This is reasonable, since compiler can't got through all functions and infer type for each possible function and possible argument type.
You can help it by using type annotations
julia> b(x::T) where T = dict["a1"](x)::T
b (generic function with 1 method)
julia> @code_warntype b(1.0)
MethodInstance for b(::Float64)
from b(x::T) where T in Main at REPL[15]:1
Static Parameters
T = Float64
Arguments
#self#::Core.Const(b)
x::Float64
Body::Float64
1 ─ %1 = Base.getindex(Main.dict, "a1")::Function
│ %2 = (%1)(x)::Any
│ %3 = Core.typeassert(%2, $(Expr(:static_parameter, 1)))::Float64
└── return %3
So compiler will add runtime check of the return type, and it will slow down tight loops, but at least it remove type instability.
Thanks for the replies. @Frederik Banning is right that I just generate this dict once, so a NamedTuple would work equally well. And @Andrey Oskin is right that type annotations do help in my case.
Instead of using wrapper functions, could macros help? I was thinking something like this:
a1(x) = x^2
a2(x) = x^3
const ntup = @NamedTuple{a1,a2}([a1,a2])
b1(x) = ntup.a1(x)
macro b2(x)
return :(ntup.a2($x))
end
println(@b2(1.0))
From what I can tell there's no type instability or performance hit in b2
vs b1
.
Then you need the same amount of macros as you have original functions.
In this case why not just use functions themselves? This macro is more or less equivalent to
const b2 = a2
It's just aliasing.
If your functions are known at compile time, you gain nothing.
And if they are unknown in compile time, you can't use macros.
Ah thanks. How many functions I'll have and their names are known at compile time, the only thing generated at runtime are the exact expressions inside the functions.
Now that I think more about it I guess I was also looking for a cost-free way to alias a more complicated function call. In practice to call one of my functions directly is (using a dict) like this funcs["H"](sol(t)...,vals...)
where vals
are global constants and sol
is an ODE solution object. I had written H(t) = funcs["H"](sol(t)...,vals...)
to save keystrokes later on, but that introduced the type instability we've been looking at. So it sounds like a macro would remove the type instability from the second step, since it's just an alias. And hopefully a NamedTuple removes some of the type instability from the first step. This way of using the ODESolution
object actually has some as well, but type annotations might help.
If it's just an alias, I would recommend to just use const H = my_long_and_hard_to_type_function_name
and not introduce extra mechanisms/macros that you have to fight with later.
It's recommended way, I suppose.
It's the same as write const Float64Q = Union{Nothing, Float64}
to deal with long type names.
As long as you do not export them it's fine.
So just to be clear, you're suggesting something like const H(t) = funcs["H"](sol(t)...,vals...)
?
Ahhh....
No, I don't :-) In this case you can just define H(t) = <initial_long_and_hard_to_type_function_name>(sol(t)..., vals...)
const
is needed when you are aliasing function names with the same signature. But here you create new function, so no const
is needed.
But main idea is to avoid funcs["H"]
in the first place if possible.
You can write macro here, btw, which would do something like
@bind alias long_and_hard_to_type_function_name
and it will execute the line that I wrote before.
Ah I see. But H(t) = <initial_long_and_hard_to_type_function_name>(sol(t)..., vals...)
is what I currently have and what seems to introduce type instability. If I type the whole thing out, I get less. I guess there's always @bind H funcs["H"]
to save a few keystrokes.
Last updated: Nov 06 2024 at 04:40 UTC