Performance optimisation: Specialised AbstactDict for VarNames?

[mhauru]

We have a lot of structures where we use Dict{T<:VarName} or sometimes OrderedDict{T<:VarName}. Since these get used in the deep internals of all VarInfo types (except SVI with NamedTuples, but that one has its own problems), and since Dicts are quite heavy and slow, I've been wondering if we could optimise this.

I'm imagining writing our own AbstractDict type that is hardcoded to have VarNames as keys. I'm not sure how to best do it, but one idea I have is that if we define an ordering on VarNames, which should be relatively natural to do (@penelopeysm tells me that she may already have code for this), we could probably exploit that by doing binary searches on a balanced binary tree or something, instead of hashes and hashmaps. Even if we stick with a hashmap, other optimisations might be available, e.g. a custom hash function. Maybe we could also use a VarNamedTuple (#900) for this.

If this gets implemented it should go in AbstractPPL, but opening the issue here since this is where it would get used.

[penelopeysm]

TuringLang/AbstractPPL.jl#142

[penelopeysm]

we could probably exploit that by doing binary searches on a balanced binary tree or something, instead of hashes and hashmaps

As per the DataStructures.jl docs (https://juliacollections.github.io/DataStructures.jl/latest/sorted_containers/#Sorted-Containers)

The containers internally use a 2-3 tree, which is a kind of balanced tree and is described in data structure textbooks. Internally, one Vector is used to store key/data pairs (the leaves of the tree) while a second holds the tree structure

but ....

using AbstractPPL, DataStructures, Chairmarks

# also copy over and run the code from https://github.com/TuringLang/AbstractPPL.jl/issues/142

pairs = [
    @varname(a) => 1.0,
    @varname(b) => 2.0,
    @varname(c) => randn(5),
    @varname(x[1]) => 4.0,
    @varname(x[2]) => ["a", "b"],
    @varname(x.z) => (a = 1, b = 2),
    @varname(y.z) => 5.0
]
d = Dict(pairs)
sd = SortedDict(pairs)

julia> @be d[@varname(x[1])]
Benchmark: 3332 samples with 1627 evaluations
 min    17.312 ns
 median 17.363 ns
 mean   17.495 ns
 max    28.734 ns

julia> @be sd[@varname(x[1])]
Benchmark: 3034 samples with 35 evaluations
 min    808.343 ns (5 allocs: 96 bytes)
 median 817.857 ns (5 allocs: 96 bytes)
 mean   826.809 ns (5 allocs: 96 bytes)
 max    1.558 μs (5 allocs: 96 bytes)

julia> @be d[@varname(c)]
Benchmark: 3437 samples with 1935 evaluations
 min    13.975 ns
 median 14.040 ns
 mean   14.243 ns
 max    29.436 ns

julia> @be sd[@varname(c)]
Benchmark: 3149 samples with 67 evaluations
 min    432.209 ns (1 allocs: 32 bytes)
 median 436.567 ns (1 allocs: 32 bytes)
 mean   443.247 ns (1 allocs: 32 bytes)
 max    828.358 ns (1 allocs: 32 bytes)

[penelopeysm]

The NT + Dict combination does make things better for identity varnames. So I wonder what happens if we shove as many VarNames as we can into the NT part, while leaving the Dict part available for really annoying ones like x[1].a.

It could be a 'progressive enhancement' thing, where we start out with a NamedTuple + DIct and then slowly expand the range of VarNames it can handle well. At the end of the day it just needs to implement getindex plus some other AbstractDict functions, the internals can be quite fluid.

using AbstractPPL: VarName, @varname
using Chairmarks

struct NTAndDict{N<:NamedTuple,D<:AbstractDict{<:VarName}}
    nt::N
    dict::D
end
Base.getindex(a::NTAndDict, k::VarName{sym,typeof(identity)}) where {sym} = a.nt[sym]
Base.getindex(a::NTAndDict, k::VarName) = Base.getindex(a.dict, k)

d = Dict{VarName,Any}(
    @varname(a) => 1.0,
    @varname(b) => 2.0,
    @varname(c) => zeros(5),
    @varname(x[1]) => 4.0,
    @varname(x[2]) => ["a", "b"],
    @varname(x.z) => (a = 1, b = 2),
    @varname(y.z) => 5.0
)
nt = (a = 1.0, b = 2.0, c = zeros(5))
dremaining = Dict{VarName,Any}(
    @varname(x[1]) => 4.0,
    @varname(x[2]) => ["a", "b"],
    @varname(x.z) => (a = 1, b = 2),
    @varname(y.z) => 5.0
)
nd = NTAndDict(nt, dremaining)

julia> @be d[@varname(x[1])]
Benchmark: 3404 samples with 1634 evaluations
 min    16.600 ns
 median 16.779 ns
 mean   16.943 ns
 max    32.334 ns

julia> @be nd[@varname(x[1])]
Benchmark: 3322 samples with 1590 evaluations
 min    17.505 ns
 median 17.715 ns
 mean   17.921 ns
 max    33.989 ns

julia> @be d[@varname(c)]
Benchmark: 3448 samples with 1891 evaluations
 min    14.322 ns
 median 14.389 ns
 mean   14.519 ns
 max    25.427 ns

julia> @be nd[@varname(c)]
Benchmark: 3564 samples with 2667 evaluations
 min    9.889 ns
 median 9.921 ns
 mean   10.017 ns
 max    17.451 ns

It's not just getindex that's faster ofc, it will probably have effects on downstream code since it's also type stable.

julia> using Test; @inferred nd[@varname(c)]
5-element Vector{Float64}:
 0.0
 0.0
 0.0
 0.0
 0.0

julia> using Test; @inferred d[@varname(c)]
ERROR: return type Vector{Float64} does not match inferred return type Any
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] top-level scope
   @ REPL[19]:1

[penelopeysm]

Here's something that's even faster at run time (getindex is 2 ns and 0 allocs, and fully type stable), but has the potential to generate a lot of compilation: https://github.com/penelopeysm/EnamRav.jl/

(It's really not much more than a gist, I just made it a package so that I could use Revise on it.)

I tried to look into why NamedTuples were so darn fast. I believe this is mostly because, for the most part, an NT is basically an anonymous C struct (https://github.com/JuliaLang/julia/blob/c0d34aa7ba29e2ba642f3ab158dc67f2be0fc19a/src/builtins.c#L1031). When you do nt.x or nt[:x] it just figures out how many bytes along is the field called x, reaches into the C struct, and reads out the value.

I think that means we don't really have any hope of adapting NamedTuple to use non-Symbol keys (because field names of structs have to be symbols). So I tried to add one level of indirection, which should only be observed at compile time in the generated function, and seems to work quite well at runtime:

• store a NamedTuple which maps symbols to values;
• also store a mapping from 'VarName's to symbols.

I don't claim that this is the final design that we should use, just trying to explore what the possible options are.

Edit: I found this design struggles a lot with loops over varnames like x[i], partly because all of them will have different types, and partly because i is value-level, so the compiler can't figure out the type of @vn(x[i]), and so it can't figure out the type of vnm[@vn(x[i])]. I added an example of this to the repo. Ah, well, you win some you lose some, this seems to be the price of having full type stability with index lenses.