Nothing and type stability in kwargs

Hi, so this is a fairly general question, but is it possible to use Nothing without type stability related performance hits? A specific example I have now is a simple simulator where I want to implement multiple optional exit conditions, say:

function run_sim(...; exit_steps=nothing, exit_time=nothing)
    exit = false
    sim_steps = 0
    sim_time = 0.0
    while !exit
        # do some simulation step, e.g.
        sim_time += rand()
        sim_steps += 1
        # check exit conditions
        if !isnothing(exit_time)
            if sim_time >= exit_time
                @info "Exited due to max simulation time met"
                exit = true
            end
        end
        if !isnothing(exit_steps)
            if sim_steps >= exit_steps
                @info "Exited due to max simulation number of steps met"
                exit = true
            end
        end
    end
end

but here exit_steps and exit_time can be either nothing which means no exit condition or a value.

Now as far as I understand if they were arguments then julia could specialize on their type when the function is called I hope even eliminate the !isnothing calls? But can this be done with kwargs instead? Iā€™d prefer that as having fixed argument positions for a potentially quite long list of different conditions can be a pain to deal with.

One ā€œsolutionā€ I can think of is to use Val to force some compile time optimizations but I get the feeling that thatā€™s not really a very good way to deal with problems. Another thing I can think of is making a struct which contains all the conditions together and then pass that as an argument (more manageable) but if I understand things correctly that struct would still need to be heavily parametrized to carry all that type information. I suspect there may be a better solution to this.

The simplest solution is to split into one user facing function with keyword arguments and an implementation function with positional arguments.

function run(args...; exit_steps=nothing, exit_time=nothing)
    _run(args..., exit_steps, exit_time)
end

function _run(..., exit_steps, exit_time)
    ...
end

The drawback is obviously that this is annoying boilerplate stuff.

2 Likes

There is an even simpler solution!

function run(x, y; exit_steps=typemax(Int64), exit_time=typemax(Int64))
    n_steps = 0
    n_seconds = ...
    while n_steps < exit_steps && n_seconds < exit_time
        n_steps += 1
        n_seconds = ...
    end
end

It cannot be used in all situations though.

3 Likes

Yes, Julia allows specializing on keyword argument type. Constant propagation should work fine, too. Otherwise, report a bug.

What makes keyword arguments different than positional arguments is that keyword arguments are not dispatched on (modulo bugs: Keyword arguments affect methods dispatch Ā· Issue #9498 Ā· JuliaLang/julia Ā· GitHub). So your code is fine.

That said, donā€™t overuse keyword arguments, as they add extra stack frames, which may cause bad performance in the presence of recursion or ugly stack traces in case of an error.

Stylistically, I prefer something like exit_steps isa Number or exit_steps isa Real.

3 Likes

Thanks for the prompt replies!

@barucden I thought of using typemax but as you say itā€™s not general.

@GunnarFarneback That is true, I wanted to avoid too much boilerplate but realistically it isnā€™t actually that much.

@nsajko This is perfect, I didnā€™t realize though it does make sense, I confused the specialization vs dispatch thing! Iā€™m glad to hear as I very much would like to rely on constant propagation and compile time optimizations. Thanks for the stack frame tip and as a side note Iā€™m somewhat mixed about the stylistic point. Your version would make it clear that it is really checking the types which is lovely, but then if I was to pass a String or something it would just silently do nothing whereas I would prefer it to fail at that point.

Actually, I hope itā€™s fine if I add one more related question (still mostly fits under the title). When managing complex kwarg structures I often find myself setting the default to nothing in the function signature and then having a block at the start of the function such as

function f(; kw1=nothing)
    if isnothing(kw1)
        kw1 = "sensible_default"
    end
end

however, if I understand the term correctly this is also a type instability isnā€™t it? Given the nature of Nothing it seems to me like a reasonable one and so I hope that this comes without a performance cost. Does it? Iā€™m not sure if Julia is smart enough to figure this out or how to test this.

1 Like

As long as the type of the argument is known to be Nothing, both the branch and the type instability should get optimized out. Additionally, even when the type is not concretely known, but it is known to subtype a small-enough type Union (less than four types, to be specific, if I remember correctly), the performance should be fine because of the ā€œunion splittingā€ optimization that the compiler is often able to do.

The usual tools are:

  • @code_typed, code_typed, @code_warntype, code_warntype, @code_llvm, code_llvm. Some of these are from the Base module, others are from the InteractiveUtils standard library package. All are available and loaded in an interactive REPL session, while a noninteractive REPL session might require loading InteractiveUtils first.
  • @descend, descend from the Cthulhu.jl package
  • @report_call, report_call, @report_opt, report_opt from the JET.jl package

Each of these has mostly the same interface, except that the macro versions try to be more user friendly by letting the user provide values instead of explicit types.

An example demonstrating the branch getting optimized out when the argument type Nothing is specialized on:

julia> function f(a = nothing)
           if isnothing(a)
               a = 3
           end
           a
       end
f (generic function with 2 methods)

julia> code_typed(f, Tuple{})  # same as `@code_typed f()`
1-element Vector{Any}:
 CodeInfo(
1 ā”€     return 3
) => Int64

julia> function g(; a = nothing)
           if isnothing(a)
               a = 3
           end
           a
       end
g (generic function with 1 method)

julia> code_typed(g, Tuple{})  # same as `@code_typed g()`
1-element Vector{Any}:
 CodeInfo(
1 ā”€     return 3
) => Int64

The above also shows that thereā€™s no difference between keyword (as in g) and positional (as in f) arguments.

An example to demonstrate the union-splitting optimization:

julia> function f(collection)
           s = 0
           for e āˆˆ collection
               s += if e isa Number
                   e
               else
                   0
               end
           end
           s
       end
f (generic function with 1 method)

julia> code_typed(f, Tuple{Vector{Union{Nothing,Int}}})  # same as `@code_typed f(Union{Nothing,Int}[])`
1-element Vector{Any}:
 CodeInfo(
[...]
) => Int64

FWIW one could use a dispatch constraint on the argument of the method to guard against that. Itā€™s preferable to throw before the method gets called (due to dispatch) rather than in the body of the method, as that makes stack traces friendlier, and also might play better with reflection tools such as isapplicable and hasmethod. Using reflection is discouraged, though.

2 Likes

Perfect! Iā€™m glad to hear and thanks for the in depth answers, examples and tips!! Iā€™m happy that I mostly donā€™t need to change how I do things and I have started using Cthulhu which has been a great improvement and also shows some of the optimized statements as Core.Const(true) etc.

1 Like