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 concreteVal-typed argument in positionidx, then returns all parameter values for theVal-typed argument as a tuple. Optionally,ptypecan be specified to filter parameter values that are instances ofptype.This function is statically compiled, and will automatically be recompiled whenever a new method of
fis 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), returnstrueif there exists a matching method with aVal-typed argument in positionidxwith parameterparamand parameter typeptype.
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!)
