Stream: helpdesk (published)

Topic: Type uncertainty from wrapper functions?

Robbie Rosati (Mar 13 2021 at 01:51):

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?

Frederik Banning (Mar 15 2021 at 10:56):

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.)

Andrey Oskin (Mar 15 2021 at 11:27):

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 is will slow down tight loops, but at least it remove type instability.

Robbie Rosati (Mar 15 2021 at 14:50):

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.

Andrey Oskin (Mar 15 2021 at 14:52):

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

Andrey Oskin (Mar 15 2021 at 14:52):

It's just aliasing.

Andrey Oskin (Mar 15 2021 at 14:52):

If your functions are known at compile time, you gain nothing.
And if they are unknown in compile time, you can't use macros.

Robbie Rosati (Mar 15 2021 at 15:12):

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.

Andrey Oskin (Mar 15 2021 at 15:16):

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.

Andrey Oskin (Mar 15 2021 at 15:17):

It's recommended way, I suppose.

Andrey Oskin (Mar 15 2021 at 15:18):

It's the same as write const Float64Q = Union{Nothing, Float64} to deal with long type names.

Andrey Oskin (Mar 15 2021 at 15:18):

As long as you do not export them it's fine.

Robbie Rosati (Mar 15 2021 at 15:20):

So just to be clear, you're suggesting something like const H(t) = funcs["H"](sol(t)...,vals...)?

Andrey Oskin (Mar 15 2021 at 15:27):

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.

Andrey Oskin (Mar 15 2021 at 15:28):

But main idea is to avoid funcs["H"] in the first place if possible.

Andrey Oskin (Mar 15 2021 at 15:29):

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.

Robbie Rosati (Mar 15 2021 at 15:44):

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: Oct 02 2023 at 04:34 UTC