Optimizing access pattern for static and deterministic models

[mohamed82008]

So after type stability issues are taken care of, one more optimization that can be worked on is how to best access fields in the VarInfo struct? The main set of optimizations I can think of make the assumption that:

1. The model is static, i.e. has a fixed number of parameters, and
2. The model has deterministic access pattern for the random variables, so any 2 variables are always evaluated in the same order.

This unlocks a whole set of possible optimizations for most if not all samplers, definitely HMC:

1. We can use StaticArrays.MArray for the fields of TypedVarInfo since it won't need to grow in size.
2. We can avoid the need for the idcs dictionary which looks up the variable by its VarName. Instead, we can sort the ranges and vals vectors by the access order of the parameters and only keep a counter for each symbol that gives us the index in ranges which gives us the index range in vals. This counter can be incremented inline in the model when expanding the @model macro. This has 2 effects. Firstly, we avoid the overhead of looking up the index in idcs and secondly we make sure the random variables are accessed in a contiguous manner in memory. This will be faster since TypedVarInfo is typically heap-allocated.

With the 2 optimizations in-place and a whole bunch of inlining, it may be possible for the Julia compiler to avoid allocating any memory on the heap for TypedVarInfo! The Julia compiler is good at eliding allocations when mutable objects only live within a single function body. This means if we make sure all the functions which use vi::TypedVarInfo are properly inlined, the compiler may be able to elide these allocations entirely. Unfortunately, this is not true for Base.Array since these are implemented in C so they have special semantics.

This optimization promises significant speedups for small and medium-sized, static, determinisitic models.

[cpfiffer]

I think the counter idea is probably a good idea in theory, but I wonder how well it'll hold up in practice. Given our push towards more community inclusion in terms of package design, is it possible to make the VarInfo stuff extremely robust given an externally produced sampler? The fixed access pattern may cause troubles for some people unaware of what might be viewed as strange behavior.

As for the fields of the VarInfo, how does using a static array provide a performance improvement? Aren't the fields themselves fairly quick to read/write? Regardless, I think the dictionary access has to go. It's kind of a band-aid that really shouldn't be in the middle of a performant system.

[mohamed82008]

Well, the fixed access pattern is a property of the models not the sampler. So ideally the sampler choice should not interfere with that. This can be triggered with an opt-in strategy. For instance, we can allow the passing of different keyword arguments to the @model macro, e.g. @model .... static = true, det = true. We can also make some diagnostic test function to make sure the the model is truly static and deterministic. This can compare the run of the model with the static VarInfo and counter with that of a traditional VarInfo. We can do this by running the model multiple times, but in reality we can only be as sure as the user is since the possibility of an untriggered randomness is hard to eliminate. That said, for the advanced users, which we seem to have a fair number of, this is a nice feature to have for small and medium sized models.

As for the fields of the VarInfo, how does using a static array provide a performance improvement?

In general, the Julia compiler isn't too liberal with optimizations involving Base.Array as it is with StaticArrays. The power of static arrays comes from the compiler being able to more heavily optimize it. This includes loop unrolling, eliding allocations, faster linear algebra, among others. If you let this propagate through the logpdf, AD and the entire model, those improvements can build up and become significant. So it is not about the read and write per se, it is about what we do with what we read after, e.g. https://github.com/TuringLang/Turing.jl/blob/master/src/inference/hmc.jl#L598.

The dictionary access is hard to get rid of in the most general case where variables can pop up at any order and any point in the sampling. You need a way to keep track of where everything is, and idcs plays this critical role.

[xukai92]

I think we might not be able to use order only, as we need to access by name for the Gibbs interface. But if we aim a specialized VarInfo for HMC and PG only, this proposal makes a lot of sense (and I also think this is the right direction).

Also, another specialized VarInfo can be created for PG that stores r.v.s with Julia types instead of vectorized version. Then we can create different types of VarInfo that this is necessary for different inference algorithm, and get the best performance while keeps the flexibility! (Ofc this could also be triggered by something like @model .... static = true, det = true or @staticmodel as you mentioned.)

[yebai]

Closed in favour of #309