[ANN] KeywordCalls.jl

I’m happy to announce KeywordCalls.jl. Registration is pending.

This problem comes up in Soss.jl and MeasureTheory.jl, so I hope making it a general-use package might also be helpful to others.

From the README

KeywordCalls allows declarations

@kwcall f(c,b,a)

with the result that

  • Calls to f(x,y,z) (without keywords) are dispatched to f(c=x, b=y, a=z)

  • Call with keywords, e.g. f(a=z, b=y, c=x) are put back in the declared “preferred ordering”, again dispatching to f(c=x, b=y, a=z)

For example,

# Define a function on a NamedTuple, using your preferred ordering
julia> f(nt::NamedTuple{(:c,:b,:a)}) = nt.a^3 + nt.b^2 + nt.c
f (generic function with 1 method)

# Declare f to use KeywordCalls
julia> @kwcall f(c,b,a)
f (generic function with 4 methods)

# Here are the 4 methods
julia> methods(f)
# 4 methods for generic function "f":
[1] f(; kwargs...) in Main at /home/chad/git/KeywordCalls/src/KeywordCalls.jl:52
[2] f(nt::NamedTuple{(:c, :b, :a), T} where T<:Tuple) in Main at REPL[2]:1
[3] f(nt::NamedTuple) in Main at /home/chad/git/KeywordCalls/src/KeywordCalls.jl:50
[4] f(c, b, a) in Main at /home/chad/git/KeywordCalls/src/KeywordCalls.jl:54

# Now other orderings work too. Here's passing a `NamedTuple`:
julia> f((a=1,b=2,c=3))
8

# Or kwargs...
julia> f(a=1,b=2,c=3)
8

# Unnamed arguments expect the declared `(c,b,a)` ordering:
julia> f(1,2,3)
32

julia> using BenchmarkTools
# Already pretty fast
julia> @btime f((a=1,b=2,c=3))
1.172 ns (0 allocations: 0 bytes)
8

# But not yet perfect, hopefully we can find a way to shave off that last nanosecond :)
julia> @btime f((c=3,b=2,a=1))
0.020 ns (0 allocations: 0 bytes)
8

Multiple declarations are allowed, as long as the set of names is distinct for each declaration of a given function.

Most of the heavy lifting is done using NestedTuples.jl and GeneralizedGenerated.jl. By taking advantage of type-level information for named tuples, we can make all of this work at compile time.

Limitations

KeywordCalls tries to push as much of the work as possible to the compiler, to make repeated run-time calls fast. But there’s no free lunch, you either pay now or pay later.

If you’d rather avoid the compilation time (at the cost of some runtime overhead), you should try KeywordDispatch.jl.

16 Likes

Just confirming that there actually is some overhead (hard to tell sometimes)

julia> function g(f, names)
           s = 0.0
           for j in 1:1000
               nt = NamedTuple{names}(randn(3)...)
               s += f(nt)
           end
           return s
           end
g (generic function with 1 method)

julia> using BenchmarkTools

julia> @btime g($f, (:a,:b,:c))
  1.091 ms (12000 allocations: 437.50 KiB)
1004.289966028431

julia> @btime g($f, (:c,:b,:a))
  930.914 μs (12000 allocations: 437.50 KiB)
784.432252203218
1 Like

I’d really love be able to do everything at compile time. The problem is that, while the keys of a NamedTuple are known at compile time, we don’t have the values until runtime. That shows up here:

@gg function kwcall(::F, nt::NamedTuple{N}) where {F,N}
    f = F.instance
    π = Tuple(kwcallperm(f, N))
    Nπ = Tuple((N[p] for p in π))
    quote
        v = values(nt)
        valind(n) = @inbounds v[n]
        $f(NamedTuple{$Nπ}(Tuple(valind.($π))))
    end
end

Here π is the permutation to get us from the call ordering to the “preferred” ordering (the one from the @kwcall declaration). I think all of the overhead is from valind.($π), though it could also have to do with inlining. Please let me know if you see a way to squash this!

As some more benchmarks demonstrate conclusively, benchmarking is hard.

I’ll put details about this in the README

2 Likes

We had lots of activity on this package yesterday, so I’ll give an update here.

KeywordCalls was working great for small tests, even across modules. But I started seeing problems when using it from MeasureTheory.jl, the use case that was the original motivation for the KeywordCalls package.

When I posted the problem on Zulip, @simeonschaub replied, “Maintaining a global cache populated by macros like this is generally something that doesn’t play very nicely with precompilation.”, and suggested it might be possible to do this just using generated functions.

@rfourquet very quickly came up with a nice workaround, including a way to include precompilation in a test. Rafael’s PR is here, and is definitely worth checking out to see some ways of working around these things.

Shortly after this, Simeon came up with a much shorter implementation of the entire package. It had some very minor issues that we worked through together, but the code is essentially his.

And it’s only 15 lines, and has no dependencies!! Here’s the whole thing:

module KeywordCalls

export @kwcall

@generated _sort(nt::NamedTuple{K}) where {K} = :(NamedTuple{($(QuoteNode.(sort(collect(K)))...),)}(nt))

function _call_in_default_order end

# Thanks to @simeonschaub for this implementation 
macro kwcall(ex)
    @assert Meta.isexpr(ex, :call)
    f, args... = ex.args
    f, args, sorted_args = esc(f), QuoteNode.(args), QuoteNode.(sort(args))
    return quote
        KeywordCalls._call_in_default_order(::typeof($f), nt::NamedTuple{($(sorted_args...),)}) = $f(NamedTuple{($(args...),)}(nt))
        $f(nt::NamedTuple) = KeywordCalls._call_in_default_order($f, _sort(nt))
        $f(; kw...) = $f(NamedTuple(kw))
    end
end

end

Now there are no precompilation issues, runtime performance is exactly the same (basically free), and the implementation is easier to understand. I haven’t measured the difference in compile time, but I’d guess it’s lower.

I think the best way forward is to use Simeon’s solution with Rafael’s tests. I’ve made a PR for this here. It needs some small updates to work with Julia 1.4, but other than that I think it’s good to go.

Thanks to Simeon and Rafael for all the help with this!

5 Likes

New PR for this:

    @kwstruct Foo(b,a,c)

Equivalent to `@kwcall Foo(b,a,c)` plus a definition

    Foo(nt::NamedTuple{(:b, :a, :c), T}) where {T} = Foo{(:b, :a, :c), T}(nt)

Note that this assumes existence of a `Foo` struct of the form

    Foo{N,T} [<: SomeAbstractTypeIfYouLike]
        someFieldName :: NamedTuple{N,T}
    end

Link:

1 Like

Interesting package! I wrote a simple Fibonacci benchmark to compare the function call overhead:

using KeywordDispatch, KeywordCalls

function fib(n,f0)
    if n <= 1
        return f0
    end
    fib(n-1,f0) + fib(n-2,f0)
end

function fib_tup((n,f0))
    if n <= 1
        return f0
    end
    fib_tup((n-1,f0)) + fib_tup((n-2,f0))
end

function fib_kw(;n,f0)
    if n <= 1
        return f0
    end
    fib_kw(;n=n-1,f0=f0) + fib_kw(;n=n-2,f0=f0)
end


@kwdispatch fib_kd()
@kwmethod function fib_kd(;n,f0)
    if n <= 1
        return f0
    end
    fib_kd(n=n-1,f0=f0) + fib_kd(n=n-2,f0=f0)
end

function fib_kc((n,f0)::NamedTuple{(:n, :f0)})
    if n <= 1
        return f0
    end
    fib_kc(n=n-1,f0=f0) + fib_kc(n=n-2,f0=f0)
end
@kwcall fib_kc(a,b)

After warmup:

julia> @time fib(30,1)
  0.006100 seconds
1346269

julia> @time fib_tup((30,1))
  0.007237 seconds
1346269

julia> @time fib_kw(n=30,f0=1)
  0.011220 seconds
1346269

julia> @time fib_kd(n=30,f0=1)
  0.318386 seconds
1346269

julia> @time fib_kc(n=30,f0=1)
  0.022082 seconds
1346269

So it appears there is some overhead to use of keyword functions, but KeywordCalls is within 3x. I’ll need to look into why KeywordDispatch is so much worse.

2 Likes

Great idea, thanks @simonbyrne!

Your @kwcall fib_kc(a,b) should be @kwcall fib_kc(n, f0). This defines the “preferred ordering” for a given set of names. This would also be a great way to test the effect of using the wrong ordering.

With that, it looks like there’s still some small overhead in calling things like fib_kc(n=n-1,f0=f0). I wasn’t aware of this, so this is really helpful.

BTW, I had originally looked at KeywordDispatch as a way to solve this problem. I would have made a PR for this, but at the time it looked like I would need some extra dependencies (mostly GeneralizedGenerated.jl). I appreciate that your work got me thinking about this :slight_smile:

1 Like