stan-dev / stan-dev/math

Use compiler attributes in stack allocator for better inlining

Open
#2,303 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
C++
Stars
839
Forks
220
Avg merge
2d 4h
Merged PRs (30d)
14

Description

Description

We can get much nicer inlining when retrieving memory from the stack allocator by using some compiler attributes. When we call alloc(), the only time we go down the path to allocate new memory is on the first iteration. Every iteration after that we go down the path that just grabs the pre-allocated memory. We can use the same cold path trick we use for the errors as well as some strong inlining for our hot path

The godbolt below shows the current stack_alloc vs one with some of the compiler attributes

https://godbolt.org/z/3TzEEd

Note that alloc_array() is tagged no-inline to replicate what optimization would look like in Stan. Since we call alloc array through a thread local static pointer the compiler is never going to inline alloc_array(). Essentially unless there are some very specific conditions to a pointer, going through one almost always turns off inlining.

doing ctrl+f for stack_alloc::alloc_array<double> shows some niceties

  • The compiler inlines everything inside of alloc_array() and with the space it saves from moving the allocating code to a cold path is able to inline / elide some extra instructions
  • When we start going down the path that we can allocate new memory (you can right click the colored code in alloc() for either source to see the asm there) the current code does a ton of jumping around when it thinks it might have to allocate new memory. Placing the alloc in an immediately invoked lambda simplifies the branch prediction logic so the instruction cache has to deal with less stuff.

There's a few other things here that we should probably discuss whether they are safe

  1. I removed all exception logic from the stack allocator. tmk any exceptions here would terminate the program with a bad_alloc() exception and personally I don't see the difference between the program crashing and unrecoverably terminating.

  2. The new version of stack_alloc uses Eigen::internal::aligned_malloc(). This can be nice since all of our memory and Eigen's memory will have the same alignment.

Expected Output

Overall this should produce nicer asm and give us a lil' speed bump. Though I still need to do testing to confirm this is true. All that we would need here is just some google benchmark that just calls alloc_array<>() and recover_memory() bunch of times

Current Version:

v3.4.0

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the stack allocator's alloc() and alloc_array<> paths, then compare the current and attribute-based implementations in the linked Godbolt example. Benchmark repeated alloc_array<> and recover_memory() calls, and confirm that the resulting assembly and measured performance improve without introducing unsafe exception or alignment behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
performance
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.