[ANN] ValSplit.jl: Compile away dynamic dispatch on Val-typed arguments

Ever wanted to use methods with Val-typed arguments, but avoid dynamic dispatch? Here’s a little package I worked on over the past week that allows you to do just that!

ValSplit.jl provides the @valsplit macro, which compiles away dynamic dispatch over methods with Val-typed arguments by “Val-splitting” (similar to union splitting). By annotating a function definition with @valsplit and choosing arguments to split upon, the resulting function will compile to a switch statement over all Val parameters associated with the chosen arguments. (Requires Julia 1.3 and above to avoid #265-like issues)

Example

Suppose we have a function soundof that takes in a Val-typed argument, and returns how an animal sounds:

soundof(animal::Val{:dog}) = "woof"
soundof(animal::Val{:cat}) = "nyan"

We might want a version of soundof that takes in Symbol values directly, and hence define:

soundof(animal::Symbol) = soundof(Val(animal))

However, when using soundof(animal::Symbol) in another function, dynamic dispatch might occur if Julia cannot infer the value of the argument animal at compile time, resulting in considerable slowdowns.

Using @valsplit, we can avoid this issue by compiling away the dispatch logic as a switch statement. We do this simply by annotating our method definition with @valsplit, and annotating each argument x::T we want to switch upon as Val(x::T):

@valsplit function soundof(Val(animal::Symbol))
    error("Sound not defined for animal: \$animal")
end

The resulting function effectively compiles to the following switch statement, where the original method body is used as the default branch:

function soundof(animal::Symbol)
    if animal == :dog
        return "woof"
    elseif animal == :cat
        return "nyan"
    else
        error("Sound not defined for animal: \$animal")
    end
end

However, unlike a manually-written switch statement, @valsplit-defined functions will automatically recompile when new methods are added. For example, if we add the method:

soundof(animal::Val{:human}) = "meh"

Then soundof(animal::Symbol) will recompile to a switch statement with an additional branch:

function soundof(animal::Symbol)
    if animal == :dog
        return "woof"
    elseif animal == :cat
        return "nyan"
    elseif animal == :human
        return "meh"
    else
        error("Sound not defined for animal: \$animal")
    end
end

As such, @valsplit-annotated functions preserve extensibility, while achieving the run-time performance of switch statements (or better, if constant propagation results in compile-time pruning of branches).

Motivation

The @valsplit macro is intended to address the following two issues:

  • Dynamic dispatch over Val-typed arguments is slow
  • Alternative solutions such as manually-written switch statements and global dictionaries are often insufficient for the purposes of extensibility.

Note that dynamic dispatch does not always occur: When there are a small number of values to split on (less than 4, as of Julia 1.6), the Julia compiler automatically generates a switch statement:

soundof(animal::Val{:dog}) = "woof"
soundof(animal::Val{:cat}) = "nyan"
soundof(animal::Symbol) = soundof(Val(animal))

julia> @code_typed soundof(:cat)
CodeInfo(
1 ─ %1  = invoke Main.Val(_2::Symbol)::Val{_A} where _A
│   %2  = (isa)(%1, Val{:cat})::Bool
└──       goto #3 if not %2
2 ─       goto #6
3 ─ %5  = (isa)(%1, Val{:dog})::Bool
└──       goto #5 if not %5
4 ─       goto #6
5 ─ %8  = Main.soundof(%1)::String
└──       goto #6
6 ┄ %10 = φ (#2 => "nyan", #4 => "woof", #5 => %8)::String
└──       return %10
) => String

But once more methods are defined, the Julia compiler no longer performs this optimization:

for i in 1:4
    sound = "sound $i"
    eval(:(soundof(animal::Val{Symbol(:animal, $i)}) = $sound))
end
soundof(animal::Symbol) = soundof(Val(animal))

julia> @code_typed soundof(:animal1)
CodeInfo(
1 ─ %1 = invoke Main.Val(_2::Symbol)::Val{_A} where _A
│   %2 = Main.soundof(%1)::Any
└──      return %2
) => Any

To avoid dynamic dispatch, manually switching on a set of values is the fastest in terms of both compile-time and run-time, but the set of values to switch upon cannot be extended. Global dictionaries can partially address this problem by associating values with code:

const SOUND_OF = Dict{Symbol,Function}()

woof() = "woof"
SOUND_OF[:dog] = woof

nyan() = "nyan"
SOUND_OF[:cat] = nyan

soundof(animal::Symbol) = SOUND_OF[animal]()

However, dictionary lookup times are usually slower compared to (small) switch statements. In addition, this approach runs into issues with precompilation, preventing a downstream module from adding new entries to a global dictionary defined in another module (except at run-time using the __init__ function). In other words, global dictionaries are not extensible across module boundaries.

The @valsplit macro addresses this problem because new methods can always be introduced by downstream modules, resulting in recompilation of the @valsplit annotated function. It effectively uses Julia’s method table as a global dictionary, but avoids the overhead of dynamic dispatch using the same @generated function tricks used to implement static_hasmethod in Tricks.jl.

A small benchmark is provided here. With 10 values to branch on, running Julia 1.6.1 on a Windows machine, the results of the benchmark are as follows:

Manual switch statement:
  3.275 μs (0 allocations: 0 bytes)
Global Dict{Symbol,String}:
  78.800 μs (0 allocations: 0 bytes)
Global LittleDict{Symbol,String}:
  111.600 μs (0 allocations: 0 bytes)
Dynamic dispatch:
  2.300 ms (0 allocations: 0 bytes)
Val-splitting with @valsplit:
  3.275 μs (0 allocations: 0 bytes

Utilities

ValSplit.jl provides a few other utility functions for determining whether a method with particular Val-typed argument exists.

To determine the set of all Val parameters associated with a particular argument of a particular function, use valarg_params:

valarg_params(f, types::Type{<:Tuple}, idx::Int, ptype::Type=Any)

Given a method signature (f, types), finds all matching methods with a concrete Val-typed argument in position idx, then returns all parameter values for the Val-typed argument as a tuple. Optionally, ptype can be specified to filter parameter values that are instances of ptype.

This function is statically compiled, and will automatically be recompiled whenever a new method of f is defined.

To determine whether a particular argument of a particular function has a specific Val parameter, use valarg_has_param:

valarg_has_param(f, types::Type{<:Tuple}, param, idx::Int, ptype::Type=Any)

Given a method signature (f, types), returns true if there exists a matching method with a Val-typed argument in position idx with parameter param and parameter type ptype.


Hope the Julia community enjoys using this package, especially developers who want to make their libraries more extensible! I’ll personally be using it to refactor the PDDL.jl parser, interpreter, and compiler so that it is faster and more extensible, and hope it comes in handy for others too.

(Also s/o to @oxinabox for writing ExprTools.jl and Tricks.jl, without which this package would have been considerably more painstaking to write, as well as some helpful suggestions for benchmarks!)

43 Likes

Niiiiice. I’ve resorted to manually hard-coding the switches in the past: MathOptInterface.jl/read.jl at 03640abedfeb1716b62ab930d091e7013abc4c73 · jump-dev/MathOptInterface.jl · GitHub

2 Likes

Now merged into the General registry! :slight_smile:

5 Likes

This is super helpful and sounds like the kind of thing that should have first class support from the language. In my understanding, this “splitting” can be done transparently by the compiler without changing semantics.

Val-dispatch makes writing parsers and handling unstructured data more elegant. IMO, we could even have syntactic sugar for these :slight_smile:

4 Likes