How do I make a dispatch table using multiple dispatch instead of Dict?

In python, I often used a dictionary-based multiple dispatch table to send arguments to functions:

dispatch = {(0,0):vrr_ss,(0,1):vrr_sp,(1,0):vrr_ps}
if a,b in dispatch: dispatch[a,b](args)

The Julia version of this is incredibly slow, I think because I can’t get the types correct:

dispatch = Dict((0,0) => vrr_ss,(0,1)=> vrr_sp, (1,0) => vrr_ps)
if (a,c) in keys(dispatch)
    return dispatch[a,c](aexpn,bexpn,cexpn,dexpn, A,B,C,D)
end

It seems that there are two ways to do this better:

  1. Figure out the dictionary type that corresponds to a mapping of pairs of integers to functions that return Matrices of Float64 (Dict{Tuple{Int64, Int64}, ???}), or
  2. Figure out how to use the standard type system to let Julia’s multiple dispatch do this for me. The Rock, Paper, Scissors blog post is suggestive, but I’m not smart enough to figure out how to take this the rest of the way.

Can anyone advise me on a way to proceed? Thanks in advance.

What problem are you trying to solve? Why do you need a table of functions indexed by pairs of integers? It’s hard to give useful advice without some context.

4 Likes

Julia naturally has multiple dispatch based on types, not values. So you could write something like this:

foo(a::Int, b::Int) = print("hi")
foo(a::AbstractFloat, b::AbstractFloat) = print("goodbye")
foo(a::Int, b::AbstractFloat) = print("what now?")


julia> foo(a::Int, b::Int) = print("hi")
foo (generic function with 1 method)

julia> foo(a::AbstractFloat, b::AbstractFloat) = print("goodbye")
foo (generic function with 2 methods)

julia> foo(a::Int, b::AbstractFloat) = print("what now?")
foo (generic function with 3 methods)

julia> foo(1, 2)
hi
julia> foo(1.0, 2.0)
goodbye
julia> foo(1, 2.0)
what now?
julia> foo(1.0, 2)
ERROR: MethodError: no method matching foo(::Float64, ::Int64)
Closest candidates are:
  foo(::Int64, ::Int64) at REPL[12]:1
  foo(::AbstractFloat, ::AbstractFloat) at REPL[13]:1
Stacktrace:
 [1] top-level scope
   @ REPL[18]:1

If you can post some more details, it’ll be easier to figure out how the dispatch you’re looking at can map to types.

1 Like

(Worrying that you’re going to be sorry you asked.)

Hokay, here’s what I’m trying to solve now. I’m working on the MolecularIntegrals.jl package, which computes electron repulsion integrals for quantum chemistry calculations. In this field, code generation is a common technique, since the recurrence relations involved lead to code that compilers can’t optimize very well. So I’d like to be able to support faster routines, either via code generation or through metaprogramming techniques. This could lead to a loop that looks like:

Do I have a faster routine for a system for angular momenta (a,c)?

  • If so, call that routine
  • If not, continue with the standard code.

The loop I posted is essentially the above question. My general code computes integrals for basis functions of arbitrary angular momenta atomic orbitals (the familiar s, p, d orbitals from introductory chemistry). But I’ve also gotten working a routine to autogenerate code customized for, e.g., s,d pairs..

What’s exciting about using Julia for this is that if I can autogenerate working code, I can perhaps also find metaprogramming solutions that obviate the need for autogenerated code. This can provide code that is relatively performant without being totally obfuscated. I’d like the next generation of graduate students who think about these programs have simple methods to learn from to figure out even better methods.

3 Likes

@stevengj I should also note that I’ve enjoyed your video on metaprogramming and am not totally naive about the problems involved. Based on some work I did yesterday, the autogenerated code is now slower than the regular Julia code, so maybe I won’t need it. But I’d like the option of doing things like this, and would like an easy way to drop code like that into the routines.

1 Like

What exactly did you find to be slow? The lookup of the pair in the list, or something else? How fast has to be the lookup relative to the time required by the function calls?

(Is it the case where many of these integrals are already computed and you are just fetching the value?)

1 Like

@lmiq I believe it’s the lookup. I’m basing this on the difference between using the dispatch table to look up function and just doing a big if/then loop:

    if mmax == 1
        return vrr_ss(aexpn,bexpn,cexpn,dexpn, A,B,C,D)
    elseif cmax == 0
        if amax == 1
            return vrr_ps(aexpn,bexpn,cexpn,dexpn, A,B,C,D)
        elseif amax == 2
            return vrr_ds(aexpn,bexpn,cexpn,dexpn, A,B,C,D)
        end
    elseif cmax == 1
        if amax == 0
(etc)

In Python, it’s thought to be much faster to write a dispatch table rather than the nested if loops, because dictionaries are fast, and many other things are slow. In Julia this does not appear to be the case, likely because many of those slow steps are faster in Julia.

3 Likes

Maybe there are two things to think about here:

  1. It that dict lookup slower in Julia than in Python? I would not expect that, if so, probably it is something to be improved and reported.

  2. The efficiency of the lookup on a dict relative to a manual branched code like that will probably depend on the data structure. For example, in the snippet you provide, it seems that from the first index you may be already decide which function to choose (if mmax==1). That is certainly much faster than searching the complete key in the dictionary.

(as a side question: how is that lookup implemented in Gamess or Gaussian?)

1 Like

I also wouldn’t expect dictionary lookup to be slower in Julia, although it’s one of the more optimized parts of python since python uses dictionaries for so many core components (like dispatch of object functions).

I can of course just write the branched code, but it feels a little inelegant, and it’s much easier to just drop in the dispatch table that uses a dictionary.

Many working quantum chemistry packages now uses Ed Valeev’s excellent libint libraries. Gaussian is commercial, and I don’t know what they use. Libint’s methods are very fast, but they gain this speed by generating and compiling totally unbranched and unrolled code customized to particular pairs and quartets of basis functions. Not sure how Libint does calls the routines. I’m wading through the code now for other reasons (trying to optimize the Boys function), and will check.

1 Like

So, the idea is to generate all (possible) functions specifically for each quartet, and the question is which would be the fastest way to dispatch on those functions, that may be hundreds of thousands, correct? If I remember from my quantum chemistry years some of these functions will simply return a precomputed value, such that the bottleneck can be the lookup itself.

You could generate custom types for each quartet and use the Julia dispatch system. Someone will probably know if that is a good idea. Or which would be the best data structure for the case.

(Yet if specific problem structure can be used, it is hard to beat a customized search)

What about haskey(dispatch,(a,c)) instead of (a,c) in keys(dispatch)?

That’s the general idea. Oddly enough, my regular code is now faster than my autogenerated code, so I may not be using that at all. But I’d like to have the option of a fast custom written routine. Plus, if I can autogenerate code, there are no doubt metaprogramming techniques and macros to explore as well.

I don’t know about hundreds of thousands of routines. HGPgen2.jl contains all pairs of s through d orbitals, and that’s only 9 routines.

1 Like

I thought that you were implementing this idea for all possible four-electron integrals.

Good point. haskey(dispatch,(a,c)) is roughly the same speed as (a,c) in keys(dispatch).

Yes, but the rate determining step for codes based on Head-Gordon and Pople’s methods are the vertical recurrence relations, which only need to compute the pairs of non-s-type angular momenta. You can then use the horizontal recurrence relations to build up the four-index arrays.

1 Like

One thing to take into consideration is that using any list of functions you can incur in run-time dispatch, which may be type unstable, and that will be much slower than the dispatch with the manual branches.

I agree with others that it may not be the best way to structure julia code, but FWIW I believe that GitHub - yuyichao/FunctionWrappers.jl is what you’re looking for the dict dispatch idea:

The idea is to have an equivalent of function pointer in C or std::function in C++11. This could be useful to avoid over specialization (Ref #13412 (comment)) and also for implementing mutable callback list/table. In these cases, it is acceptable that the callback function is not/cannot be inlined to the caller but it is also desired to avoid necessary boxing or runtime type check/dispatch.

4 Likes

I don’t understand your use case well enough to give recommendations specific to your package, but here are a few things you can play around with.

1

Code that looks like

if haskey(d, k)
    return d[k]
end

is not as fast as it could be, because two dictionary lookups are performed: one for haskey(d, k) and one for d[k]. You can get away with only one lookup by using get. Here’s a benchmark:

using BenchmarkTools

function foo(d, k)
	foreach(1:1000) do _
		if haskey(d, k)
			d[k]
		end
	end
end

function bar(d, k)
	foreach(1:1000) do _
		get(d, k, nothing)
	end
end

d = Dict(:a => 1, :b => 2)
julia> @btime foo($d, :a);
  12.001 μs (0 allocations: 0 bytes)

julia> @btime bar($d, :a);
  5.867 μs (0 allocations: 0 bytes)

So your version could look something like this:

f = get(d, k, nothing)
if f !== nothing
    return f(args...)
end

2

It’s sort of been hinted at in this thread, but I don’t think it’s been explicitly mentioned:

The type of a function is simply typeof(foo), so every function has it’s own type, even if the input and output types are ostensibly the same:

julia> foo() = 1;

julia> typeof(foo)
typeof(foo)

This means you can’t have a type stable dictionary who’s values are functions.

3

I don’t know how performant this would be, but you could potentially play around with types parameterized by Ints, or use Val types. Here’s an example:

struct Momentum{A, B} end

function Momentum(A, B)
	Momentum{A, B}()
end

repulsion(::Momentum{0, 0}) = 1
repulsion(::Momentum{1, 0}) = 2
repulsion(::Momentum{0, 1}) = 3
julia> repulsion(Momentum(0, 0))
1

julia> repulsion(Momentum(1, 0))
2

julia> repulsion(Momentum(0, 1))
3

4

Your README says that one of the goals of your package is to provide “simple, easily understandable methods”. If that’s the case, then I would advise against using macros. Macros tend to be a bit of a black box and add a whole new syntax to learn. (Although it’s often best to follow the rule of thumb that the user should be able to predict the non-macro code that a macro will expand to.)

As you’re aware from watching the metaprogramming video from @stevengj, 99% of the time you should be able to achieve performance without using code generation. Macros are just for introducing new syntax (99% of the time). As I mentioned above, I don’t really know your domain, but instead of code generation maybe you can achieve what you need with higher-order functions. :slight_smile:

5

This is just a shot in the dark, but it’s also possible that generated functions might suit your needs.

4 Likes

Excellent. Lot to look into here. Thanks for replying!

1 Like

Excellent. I’ll look into it. Thanks!