Efficient custom dict for mostly same types

General Usage
1

I have the following setup, which basically creates a linked-list-type A, wich contains some parametric type B:

struct B{T}
    val::T
end

struct A{T<:B}
    parent::A{T}
    x::T
    A{T}() where T = new()
    A(parent::A{T}, x::T) where T<:B = new{T}(parent, x)
end

A(x::T) where T<:B = A(A{T}(), x)

This works pretty well, if all types are the same:

julia> A(A(B(1.)), B(2.))
A{B{Float64}}(A{B{Float64}}(A{B{Float64}}(#undef, B{Float64}(4.0e-323)), B{Float64}(1.0)), B{Float64}(2.0))

The problem now is, that I want to allow A to contain different parametric subtypes of B, to get the following to work:

julia> A(A(B(1.)), B(2))
ERROR: MethodError: no method matching A(::A{B{Float64}}, ::B{Int64})

I could of course make A non-parametric and have it say x::B, but because most of the time, x is going to be the same type, that’s a performance penalty I am not willing to take.
What I want, is for the above to be promoted to A{B{Union{Int64,Float64}}, or even A{B{UnionAll}}, but I don’t want any x to be actually converted because I need a.x::B{Type1} === A(a, _x::B{Type2}).parent.x to still be valid.
I have tried quite a few approaches to get this to work but neither of them did, so I would be very grateful, if anyone has an idea for how to get this to work.

2

Just don’t restrict the type of the value and the parent to be the same, eg

struct A{S, T<:B}
    parent::S
    x::T
end

Note however that this will result in a long chain of nested types when the xs are heterogeneous, with the associated compilation cost. For most applications with heterogeneous elements, I imagine that either tuples would be competitive (or at least degrade better than this); for collections that have a small number of types and a lot of elements you should consider something else.

Also, unless the linked list is a programming exercise, it may not be the most efficent solution in any case.

3

Thanks for the quick response! I should probably been a bit more specific with my use case, I just thought there might be some way to do these kind of promotions that I was missing. What I actually need it for is this: https://github.com/simeonschaub/Covar.jl/blob/master/src/gradients.jl. It’s a kind of dictionary and it is very fast for what I need it to do, it just doesn’t support different Types.

As my use case mostly consists of a lot of elements and only occasionally different types, this isn’t really an efficient solution for me. In these cases it is slower than the non-parametric version and comes with a lot of overhead, because, even when the types are the same, each type has to be saved differently. My non-parametric version is ~20% slower than my parametric example, so it would be the better choice here; I just thought I could take advantage of these 20% when they are the same type. When they are of different types, a performance hit is of course to be expected.
Are these kind of promotions from a type to a union containing these types at all possible in Julia, or is it something that the language doesn’t allow?

4

Promotions are possible, but outside isbits Union optimizations, you may get non-specialized code and boxed values.

I did not read the example you linked (please try to simplify to something small and self-contained), but if you just need quick lookup for a small number of types, you can create a Tuple of Dicts, or linked lists.

Mock example using Dict (you can use your linked lists):

struct DictTuple{T}
    dicts::T
end

Base.getindex(dt::DictTuple, index) = _lookup(index, dt.dicts...)

_lookup(index) = nothing        # not found, throw an error?
_lookup(index::T, d::Dict{T}, rest...) where T = d[index]
_lookup(index, d, rest...) = _lookup(index, rest...)

dt = DictTuple((Dict(1 => 2), Dict('a' => "a fish")))

# look at these, all type stable
@code_warntype dt[1]
@code_warntype dt['a']
@code_warntype dt["not stored"]
2 Likes
5

This solution looks really promising. I’ll have to experiment a bit to find out how exactly I will implement it, but it seems that this is what I was looking for all along. Thank you very much!

6

Happy to help. The solution is a bit dense and just an MWE, so feel free to ask about modifying it.

7

You might also want to check out https://github.com/tkoolen/TypeSortedCollections.jl which is designed to efficiently operate over type-heterogeneous data with a relatively small number of possible types.

1 Like
8

One thing I have to do quite often is merge two Dicts. I basically want it to work like Base.merge for Dicts. I already have a function f that merges two Dicts of the same type, but for simplicity’s sake I will use the type B from above instead of my Dict type and define f like this:

f(x::B{T}, y::B{T}) where T = B(x.val + y.val)::B{T} #always returns same type as x and y

I then defined merge as follows, so it recursively looks for an entry in y, that has the same type as the current item from x. If it finds one, it returns f(x_i, y_j), otherwise it just returns x_i and those get added to the resulting tuple. In the end it also appends any entries from y that couldn’t be merged with x to the result and returns it.

__merge(f, unused_ys, x_i) = x_i, unused_ys
__merge(f, unused_ys, x_i::T, y_j::T, y...) where T = f(x_i, y_j), (unused_ys..., y...)
__merge(f, unused_ys, x_i, y_j, y...) = __merge(f, (unused_ys..., y_j), x_i, y...)

_merge(f, results, remaining_ys) = (results..., remaining_ys...)
function _merge(f, results, remaining_ys, x_i, x...)
    _result, _remaining_ys = __merge(f, (), x_i, remaining_ys...)
    return _merge(f, (results..., _result), _remaining_ys, x...)
end

merge(f, x, y) = _merge(f, (), y, x...)

As you can see, it has some nice recursion in it, but it seems to me that it’s not as efficient as it could be.
This for example works like it should:

julia> merge(f, (B(2.), B(1), B(0xff)), (B(3), B(3f0), B(1.)))
(B{Float64}(3.0), B{Int64}(4), B{UInt8}(0xff), B{Float32}(3.0f0))

but @code_warntype prints:

julia> @code_warntype merge(f, (B(2.), B(1), B(0xff)), (B(3), B(3f0), B(1.)))
Body::Tuple
15 1 ─ %1  = (getfield)(x, 1)::B{Float64}                                                     β”‚      
   β”‚   %2  = (getfield)(x, 2)::B{Int64}                                                       β”‚      
   β”‚   %3  = (getfield)(x, 3)::B{UInt8}                                                       β”‚      
   β”‚   %4  = (getfield)(y, 1)::B{Int64}                                                       β”‚β•»      _merge
   β”‚   %5  = (getfield)(y, 2)::B{Float32}                                                     β”‚β”‚     
   β”‚   %6  = (getfield)(y, 3)::B{Float64}                                                     β”‚β”‚     
   β”‚   %7  = (Base.getfield)(%1, :val)::Float64                                               β”‚β”‚β•»β•·β•·β•·   __merge
   β”‚   %8  = (Base.getfield)(%6, :val)::Float64                                               │││┃│││   __merge
   β”‚   %9  = (Base.add_float)(%7, %8)::Float64                                                β”‚β”‚β”‚β”‚β•»      __merge
   β”‚   %10 = %new(B{Float64}, %9)::B{Float64}                                                 β”‚β”‚β”‚β”‚β”‚β•»      f
   β”‚   %11 = (Core.tuple)(%4, %5)::Tuple{B{Int64},B{Float32}}                                 β”‚β”‚β”‚β”‚β”‚  
   β”‚   %12 = (Core.tuple)(%10)::Tuple{B{Float64}}                                             β”‚β•»      _merge
   β”‚   %13 = invoke Main._merge(_2::Function, %12::Tuple{B{Float64}}, %11::Tuple{B{Int64},B{Float32}}, %2::B{Int64}, %3::B{UInt8})::Tuple
   └──       return %13                                                                       β”‚

I might be reading that wrong, but it seems to me that it can’t infer the output type after the second call to _merge. Is there a way to give the compiler more information, for example, that the first types of the output tuple are always going to be the same as the types of tuple x? Or just to write _merge more clearly for the compiler, maybe even in an explicit loop instead of recursion?

9

I got a partial improvement with

struct B{T}
    val::T
end

f(x::B{T}, y::B{T}) where T = B{T}(x.val + y.val) #always returns same type as x and y

_merge1(f, x, y_acc) = (x, y_acc...)
_merge1(f, x::T, y_acc, y::T, y_rest...) where T = (f(x, y), y_acc..., y_rest...)
_merge1(f, x, y_acc, y, y_rest...) = _merge1(f, x, (y, y_acc...), y_rest...)
merge1(f, x, ys) = _merge1(f, x, (), ys...)

merge2(f, ::Tuple{}, ys) = ys
merge2(f::F, xs, ys) where F = merge2(f, Base.tail(xs), merge1(f, first(xs), ys))

where

julia> @code_warntype merge1(f, B(1), (B(3), B(3f0), B(1.)))
Body::Tuple{B{Int64},B{Float32},B{Float64}}
1 ─ %1 = (getfield)(ys, 1)::B{Int64}
β”‚   %2 = (getfield)(ys, 2)::B{Float32}
β”‚   %3 = (getfield)(ys, 3)::B{Float64}
β”‚   %4 = (Base.getfield)(x, :val)::Int64
β”‚   %5 = (Base.getfield)(%1, :val)::Int64
β”‚   %6 = (Base.add_int)(%4, %5)::Int64
β”‚   %7 = %new(B{Int64}, %6)::B{Int64}
β”‚   %8 = (Core.tuple)(%7, %2, %3)::Tuple{B{Int64},B{Float32},B{Float64}}
└──      return %8

julia> @code_warntype merge2(f, (B(2.), B(1), B(0xff)), (B(3), B(3f0), B(1.)))
Body::Tuple{B{UInt8},B{Float64},Vararg{Any,N} where N}
1 ─ %1  = (getfield)(xs, 2)::B{Int64}
β”‚   %2  = (getfield)(xs, 3)::B{UInt8}
β”‚   %3  = (Base.getfield)(xs, 1, true)::B{Float64}
β”‚   %4  = (getfield)(ys, 1)::B{Int64}
β”‚   %5  = (getfield)(ys, 2)::B{Float32}
β”‚   %6  = (getfield)(ys, 3)::B{Float64}
β”‚   %7  = (Base.getfield)(%3, :val)::Float64
β”‚   %8  = (Base.getfield)(%6, :val)::Float64
β”‚   %9  = (Base.add_float)(%7, %8)::Float64
β”‚   %10 = %new(B{Float64}, %9)::B{Float64}
β”‚   %11 = (Base.getfield)(%1, :val)::Int64
β”‚   %12 = (Base.getfield)(%4, :val)::Int64
β”‚   %13 = (Base.add_int)(%11, %12)::Int64
β”‚   %14 = %new(B{Int64}, %13)::B{Int64}
β”‚   %15 = (Core.tuple)(%14)::Tuple{B{Int64}}
β”‚   %16 = invoke Main._merge1(_2::Function, %2::B{UInt8}, %15::Tuple{B{Int64}}, %5::B{Float32}, %10::B{Float64})::Tuple{B{UInt8},B{Float64},Vararg{Any,N} where N}
└──       return %16

and somehow the type of the function is not inferred. Maybe someone can improve this (the code snippet above is self-contained).

10

I managed to write a version of my merge function, that has excellent type-inference, by basically manually telling the compiler that the first results of _mergeo are the same type as x:

struct B{T}
    val::T
end

f(x::B{T}, y::B{T}) where T = B{T}(x.val + y.val) #always returns same type as x and y


__merge(f, unused_ys, x_i) = x_i, unused_ys
__merge(f, unused_ys, x_i::T, y_j::T, y...) where T = f(x_i, y_j), (unused_ys..., y...)
__merge(f, unused_ys, x_i, y_j, y...) = __merge(f, (unused_ys..., y_j), x_i, y...)

@inline _mergeo(f, results, remaining_ys) = (results, remaining_ys)
@inline function _mergeo(f, results, remaining_ys, x_i, x...)
    _result, _remaining_ys = __merge(f, (), x_i, remaining_ys...)
    return _mergeo(f, (results..., _result), _remaining_ys, x...)
end

@inline function mergeo(f, x::T, y) where T
    let (res::T, rem) = _mergeo(f, (), y, x...)
        return (res..., rem...)
    end
end

As one can see, type inference works great:

julia> @code_warntype mergeo(f, (B(2.), B(1), B(0x03), B(Float16(4.))), (B(3), B(3f0), B(1.), B(0xf), B(0x0012)))
Body::Tuple{B{Float64},B{Int64},B{UInt8},B{Float16},B{Float32},B{UInt16}}
27 1 ─ %1  = (getfield)(x, 1)::B{Float64}                                               β”‚      
   β”‚   %2  = (getfield)(x, 2)::B{Int64}                                                 β”‚      
   β”‚   %3  = (getfield)(x, 3)::B{UInt8}                                                 β”‚      
   β”‚   %4  = (getfield)(x, 4)::B{Float16}                                               β”‚      
   β”‚   %5  = (getfield)(y, 1)::B{Int64}                                                 β”‚β•»      _mergeo
   β”‚   %6  = (getfield)(y, 2)::B{Float32}                                               β”‚β”‚     
   β”‚   %7  = (getfield)(y, 3)::B{Float64}                                               β”‚β”‚     
   β”‚   %8  = (getfield)(y, 4)::B{UInt8}                                                 β”‚β”‚     
   β”‚   %9  = (getfield)(y, 5)::B{UInt16}                                                β”‚β”‚     
   β”‚   %10 = (Base.getfield)(%1, :val)::Float64                                         β”‚β”‚β•»β•·β•·β•·   __merge
   β”‚   %11 = (Base.getfield)(%7, :val)::Float64                                         │││┃│││   __merge
   β”‚   %12 = (Base.add_float)(%10, %11)::Float64                                        β”‚β”‚β”‚β”‚β•»      __merge
   β”‚   %13 = %new(B{Float64}, %12)::B{Float64}                                          β”‚β”‚β”‚β”‚β”‚β•»      f
   β”‚   %14 = (Core.tuple)(%5, %6, %8, %9)::Tuple{B{Int64},B{Float32},B{UInt8},B{UInt16}}β”‚β”‚β”‚β”‚β”‚  
   β”‚   %15 = (Core.tuple)(%13)::Tuple{B{Float64}}                                       β”‚β•»      _mergeo
   β”‚   %16 = invoke Main._mergeo(_2::Function, %15::Tuple{B{Float64}}, %14::Tuple{B{Int64},B{Float32},B{UInt8},B{UInt16}}, %2::B{Int64}, %3::B{UInt8}, %4::Vararg{Any,N} where N)::Tuple{Tuple,Tuple{B{Float32},B{UInt16}}}
   β”‚   %17 = (Base.getfield)(%16, 1)::Tuple                                             β”‚β”‚β•»      indexed_iterate
   β”‚   %18 = (isa)(%17, Tuple{B{Float64},B{Int64},B{UInt8},B{Float16}})::Bool           β”‚      
   └──       goto #3 if not %18                                                         β”‚      
   2 ─ %20 = Ο€ (%17, Tuple{B{Float64},B{Int64},B{UInt8},B{Float16}})                    β”‚      
   └──       goto #4                                                                    β”‚      
   3 ─ %22 = (Base.convert)($(Expr(:static_parameter, 1)), %17)::Tuple{B{Float64},B{Int64},B{UInt8},B{Float16}}
   └──       goto #4                                                                    β”‚      
   4 β”„ %24 = Ο† (#2 => %20, #3 => %22)::Tuple{B{Float64},B{Int64},B{UInt8},B{Float16}}   β”‚      
   β”‚   %25 = (Base.getfield)(%16, 2)::Tuple{B{Float32},B{UInt16}}                       β”‚β•»      indexed_iterate
28 β”‚   %26 = (getfield)(%24, 1)::B{Float64}                                             β”‚      
   β”‚   %27 = (getfield)(%24, 2)::B{Int64}                                               β”‚      
   β”‚   %28 = (getfield)(%24, 3)::B{UInt8}                                               β”‚      
   β”‚   %29 = (getfield)(%24, 4)::B{Float16}                                             β”‚      
   β”‚   %30 = (getfield)(%25, 1)::B{Float32}                                             β”‚      
   β”‚   %31 = (getfield)(%25, 2)::B{UInt16}                                              β”‚      
   β”‚   %32 = (Core.tuple)(%26, %27, %28, %29, %30, %31)::Tuple{B{Float64},B{Int64},B{UInt8},B{Float16},B{Float32},B{UInt16}}
   └──       return %32                                                                 β”‚

I’ve also written a small function to do some microbenchmarks. You can basically specify the length of xs and ys and how many of the types should be the same. I’ve pasted it here if anyone wants to experiment.
I can’t say how representative the benchmark is for real-world applications, type-inference can probably be better leveraged in the context of a more complex program. The results were also all over the place, often times my first example was about factor 10 times faster than yours, but for other numbers of arguments, your function took in the realms of tens of nanoseconds, while mine took microseconds.
With @inline added, my type-inferable function was often a bit faster than my first attempt, but there were also times where it was a bit slower. The other functions didn’t seem to really benefit from inlining. These benchmarks made me wonder, how important full type inference here really is, or wether were kind of barking up the wrong tree by just looking at @code_warntype.
What was most important to my application, was that there is basically no performance penalty, when all entries are the same type, which all of these functions achieve, but I’ll probably do some benchmarking of these functions once I’ve fully implemented them.

1 Like