Stan SIMD & Performance

[stevebronder]

@chriselrod

Eigen already “packs” doubles while performing matrix multiplication (at least for large matrices).

Is that done in this code here? Feel like a half semester of a C++ optimization class could be spent on Eigen’s product folder haha

@chriselrod

Julia has a syntax for element-wise operations (adding a dot to a function call or operator);

I would love for stan to steal this idea. A consistent and nice shorthand for broadcasting would be awesome

@chriselrod

In calculating everything in one pass, I also tried to order computations so that values in registers can be reused rather than stored somewhere and reloaded later.

The rest of what you said above this line is extremely cool but a bit higher than my paygrade haha. But one other thing about our var type is that the base vari inside of it is 24 bytes (2 doubles and 2 virtual functions). This has always felt very not cache friendly since most modern caches are 64 bytes we can only fit 2 varis in cache at once. If we could remove those virtual functions we could fit 4 varis in each cache line. @Bob_Carpenter I think you said before you had some scheme for that?

@chriselrod

If you need dual number types for this, try to use struct-of-array type data structures

I don’t know how to do this.

@seantalts is correct, that’s exactly what I meant. If the values and adjoints can both be packed densely in their respective arrays, this would save you from copying for matrix multiplication, and allow other loops and function using vectorized/array notation (like y ~ normal(mu, sigma) when some of these are vectors) to use SIMD (again, without copying).

Bob and I have spoken about this before and I think the post here talks about adding a static matrix type for that. I never fully understood why the static matrix type would force a copy on setting

@chriselrod

I will add more documentation and benchmarks as I get closer to release. In particular, I wrote a custom allocator for InplaceDHMC to manage memory while evaluating trees / performing the leap frog steps.

I wonder if we made Stan’s stack allocator more general if we could do something like this as well? When I ran heaptrack over one of my stan programs I found it took a lot of time allocating and deallocating the point_ps types when it built the tree.

@chriselrod

In ProbabilityModels.jl, I write everything to a custom stack, simply incrementing the stack pointer each time. I know exactly what escapes (the gradient, and nothing else), so I can simply restart the stack pointer at the same place each time.

I think with a bit of weirdness and elbow grease this is feasible in our current stuff. Something something using our own stack allocator and Eigen::Map to construct matrices instead of Eigen’s allocator.

@chriselrod

I have already considering implementing a C++ version of this.

Orrr you could help us do some of this! :-)

@seantalts

I should just caveat - I haven’t profiled that many models and they can have very different bottlenecks. But yeah. I tried to use a custom allocator (I could try to find the branch) for std::vector and didn’t see much of a speedup, but didn’t investigate too hard.

I tried this as well though I also don’t think I tried hard enough. I think you need the local allocator for each object to then actually point to a global allocator that’s never destroyed. Then do the standard tracking of space used.

@seantalts

don’t think I tried anything similar for Eigen types - it’s much less clear the optimal way to switch the allocator for those (is there anything documented?).

Also goofed around with this, I think we would need to use Eigen::Map if we want to allocate our own memory for Eigen. They don’t have any macros or templates for using your own allocator (that I know of).

@seantalts

We do the same thing for Stan’s AD primitive, the var . So, if I recall correctly, we’re doing something much worse (from a SIMD perspective) than having a matrix full of structs as you were thinking; we actually have a matrix full of pointers to structs :P (each var is a pointer allocated on our custom AD stack memory pool to point to a vari struct containing a value and an adjoint). Those Eigen matrices are allocated on the heap, though the two doubles for each cell is allocated (individually) in our custom memory pool.

Yeah if we can remove the pimpl idiom we currently use in var->vari that would remove a lot of pointer chasing. I’ve tried goofing with this but can’t find a pattern that works well. Though I’m sure there is something out there to do. The main thing is dealing with the global AD tape that holds a pointer to those varis in var.

@tadej

Actually I recently starting working on that. Here is link to the github issue: https://github.com/stan-dev/math/issues/1470 . First PR is already waiting for a review. However it will take a long time to completely fix this as is needs to touch almost every single function in Stan Math.

Yes and it’s excellent! Tadej also worked on something similar to this for the elementwise functions that’s seeing a nice average speedbump of 3% on the performance benchmarks. (meta, it’s v rare to see PRs that give speedups to all of the benchmarks so that is very nice)