Repository navigation
Fix/aligned vari alloc - #3466
Open
SteveBronder wants to merge 6 commits into
Open
SteveBronder wants to merge 6 commits into
SteveBronder wants to merge 6 commits into
Conversation
reverse_pass_callback stores its captures inside the vari, which is allocated on the arena with 8 byte alignment. A capture that needs more, such as a fixed-size Eigen matrix (Matrix2d needs 32 bytes), could be misaligned and crash in vectorized code. - stack_alloc::alloc(len, alignment) returns memory aligned to `alignment`. Allocations after it stay 8 byte aligned. New blocks for alignments over 8 come from std::aligned_alloc, and blocks are only reused when the aligned allocation fits. - vari_base gains operator new(size_t, std::align_val_t), which C++17 calls for over-aligned types, and the matching operator delete. Known gaps: types needing exactly 16 bytes still use operator new(size_t) (C++17 only calls the aligned overload above __STDCPP_DEFAULT_NEW_ALIGNMENT__), and std::aligned_alloc is not available on MSVC/MinGW.
stan::math::multiply's callback captures fixed-size Eigen matrices inside the vari, so this crashed before the aligned vari allocation.
reverse_pass_callback and make_callback_vari store the callback's captures inside the vari. C++17 only calls vari_base's aligned operator new for alignments above __STDCPP_DEFAULT_NEW_ALIGNMENT__ (16 on common platforms), so a capture that needs exactly 16 bytes, e.g. a Matrix<var, 2, 1> captured by multiply, still went through operator new(size_t) and the 8 byte aligned arena. Copying the capture into the misaligned vari faulted (seen in RealSchur's Francis QR step, which multiplies fixed-size Householder vectors). Both now place the vari at memalloc_.alloc(sizeof(vari), alignof(vari)), which is unchanged for varis that need at most 8 bytes. Tests capture a Vector2d (16 bytes) and a Matrix4d (32 bytes) in each helper after leaving the arena 8 bytes past a 16 byte boundary; both segfaulted before this change.
alloc(len) and alloc(len, alignment) are now one function. Alignments of 8 or less take the old bump-pointer fast path, which folds away at compile time since callers pass a constant alignment; larger alignments round up the next location and then use the same bump and block check. The generated code for default-aligned allocations is unchanged, and over-aligned allocations no longer do separate size arithmetic. alloc_array also takes an alignment, and the class docs now state that every allocation is at least 8 byte aligned.
Add STAN_STRONG_INLINE (always_inline) to compiler_attributes.hpp and use it on the allocation fast path: stack_alloc::alloc, alloc_array, move_to_next_block, vari_base::operator new, the scalar vari_value constructors, reverse_pass_callback_vari and reverse_pass_callback. With move_to_next_block inlined the compiler can keep next_loc_ and cur_block_end_ in registers across a loop of allocations. When it is an out of line call, clang reloads and stores them through memory on every allocation. In a microbenchmark that reuses the arena (the steady state after the first gradient) clang goes from 1.34 to 0.50 ns per alloc(24), 1.55 to 0.69 ns per alloc_array, and 5.6 to 3.6 ns per new vari; gcc already inlined these and is unchanged. Add AutodiffStackStorage::push_var_stack and push_var_nochain_stack, which keep vector growth in a STAN_COLD_PATH function since the stacks are reused and almost never grow. This takes a register spill out of every vari construction with gcc (3.8 to 3.5 ns per new vari). All direct var_stack_ and var_nochain_stack_ push_back calls now use them. New arena blocks always come from aligned_malloc, so move_to_next_block no longer branches on the alignment to pick an allocator.
MSVC and MinGW do not provide std::aligned_alloc. Add STAN_ALIGNED_ALLOC(alignment, size) and STAN_ALIGNED_FREE(ptr), which use _aligned_malloc and _aligned_free on _WIN32 and std::aligned_alloc and std::free everywhere else. Memory from _aligned_malloc must be released with _aligned_free, so every arena block now comes from STAN_ALIGNED_ALLOC, including the first block that used to come from malloc, and every block is released with STAN_ALIGNED_FREE. eight_byte_aligned_malloc had no other users and is removed.
Contributor
Jenkins Console Log Machine informationDistributor ID: Ubuntu Description: Ubuntu 20.04.3 LTS Release: 20.04 Codename: focal CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Byte Order: Little Endian Address sizes: 52 bits physical, 57 bits virtual CPU(s): 192 On-line CPU(s) list: 0-191 Thread(s) per core: 2 Core(s) per socket: 48 Socket(s): 2 NUMA node(s): 2 Vendor ID: AuthenticAMD CPU family: 25 Model: 17 Model name: AMD EPYC 9474F 48-Core Processor Stepping: 1 Frequency boost: enabled CPU MHz: 1497.504 CPU max MHz: 4114.4229 CPU min MHz: 1500.0000 BogoMIPS: 7189.39 Virtualization: AMD-V L1d cache: 3 MiB L1i cache: 3 MiB L2 cache: 96 MiB L3 cache: 512 MiB NUMA node0 CPU(s): 0-47,96-143 NUMA node1 CPU(s): 48-95,144-191 Vulnerability Gather data sampling: Not affected Vulnerability Indirect target selection: Not affected Vulnerability Itlb multihit: Not affected Vulnerability L1tf: Not affected Vulnerability Mds: Not affected Vulnerability Meltdown: Not affected Vulnerability Mmio stale data: Not affected Vulnerability Old microcode: Not affected Vulnerability Reg file data sampling: Not affected Vulnerability Retbleed: Not affected Vulnerability Spec rstack overflow: Mitigation; Safe RET Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; STIBP always-on; PBRSB-eIBRS Not affected; BHI Not affected Vulnerability Srbds: Not affected Vulnerability Tsa: Mitigation; Clear CPU buffers Vulnerability Tsx async abort: Not affected Vulnerability Vmscape: Mitigation; IBPB before exit to userspace Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good amd_lbr_v2 nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba perfmon_v2 ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local user_shstk avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif x2avic v_spec_ctrl vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid overflow_recov succor smca fsrm flush_l1d debug_swap G++: g++ (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0 Copyright (C) 2019 Free Software Foundation, Inc. This is free software; see the source for copying conditions. There is NO warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. Clang: clang version 10.0.0-4ubuntu1 Target: x86_64-pc-linux-gnu Thread model: posix InstalledDir: /usr/bin |
WardBrian
reviewed
Oct 16, 2026
WardBrian
left a comment
WardBrian
left a comment
Member
There was a problem hiding this comment.
Overall this looks clear, but the cold path for growing the vector is confusing to me
Comment on lines
+138
to
+149
| STAN_STRONG_INLINE static inline void push_back_reused( | ||
| std::vector<ChainableT *> &stack, ChainableT *vi) { | ||
| if (likely(stack.size() != stack.capacity())) { | ||
| stack.push_back(vi); | ||
| } else { | ||
| grow_and_push_back(stack, vi); | ||
| } | ||
| } | ||
|
|
||
| STAN_COLD_PATH static void grow_and_push_back( | ||
| std::vector<ChainableT *> &stack, ChainableT *vi) { | ||
| stack.push_back(vi); |
Member
There was a problem hiding this comment.
How is this really an optimization? both sides of the if are just vector.push_back?
Member
There was a problem hiding this comment.
Is it relying on the compiler inlining push_back or something? If so I feel like a comment should say so
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This fixes a segfault when a reverse mode callback captures a fixed-size Eigen type, and while I was in the stack allocator I also did most of #2303.
reverse_pass_callbackandmake_callback_varistore the lambda's captures inside the vari, and the vari lives on our arena which only hands out 8 byte aligned memory. If you capture something like aMatrix<double, 4, 4>or aMatrix<var, 2, 1>, Eigen assumes the memory is aligned and uses aligned simd loads/stores, which can crash.The changes are
stack_alloc::alloc(len, alignment = 8)replacesalloc(len). Foralignment <= 8it's the same bump pointer as before. For larger alignments it rounds upnext_loc_first.alloc_array<T>(n, alignment = 8)is the array version ofalloc(...)vari_basegets an alignedoperator new/operator deleteso C++17 picks it for anything over 16 byte aligned.reverse_pass_callbackandmake_callback_varido placement new atmemalloc_.alloc(sizeof(vari_t), alignof(vari_t)), which covers the 16 byte case too.STAN_ALIGNED_ALLOCand are freed withSTAN_ALIGNED_FREE. MSVC and MinGW don't havestd::aligned_alloc, so on Windows these go to_aligned_malloc/_aligned_free. I don't have a Windows machine so CI is the first real test of that.Speeding up the allocator (#2303)
In Stan we do a ton of allocations on the first gradient and after that we are pretty much always reusing memory the arena already has. So the hot path is just bump the pointer and compare. We want to use
STAN_STRONG_INLINEandSTAN_COLD_PATHon specific functions and constructors such that the hot path we normally run is as clear and inline as possible. This is an annoying optimization problem, but one that claude is great at. So I told claude the problem and the set of functions/constructors it could add those two macros at to find a set that will give a speedup on a given google benchmark below.I wrote an MRE of develop and this PR's stack allocator to look at the asm.
https://godbolt.org/z/xYcrhY73T
clang is really conservative with inlining. On develop it leaves
move_to_next_block()as a function call, so it has to load and storenext_loc_through memory on every allocation. AddingSTAN_STRONG_INLINEtoalloc(),alloc_array(),move_to_next_block(), andvari_base::operator newlets clang keep it in a register. gcc already inlines all of this so the arena is the same speed there.I also added
push_var_stack()andpush_var_nochain_stack()toAutodiffStackStorage, which put the vector's grow path in aSTAN_COLD_PATHfunction since after the first gradient the var stack basically never grows. That takes a register spill out of every vari construction on gcc.Below is a Google Benchmark against the actual stan headers, comparing the commit this branch started from to this PR. ns per op, mean ± sd of 10 repetitions.
new varireverse_pass_callbackalloc(24)alloc_array<double>(n)Going through the hot loops with
perf annotate, here is where the time comes fromalloc/alloc_array: develop's loop loadsnext_loc_from memory, adds to it, and stores it back, so every allocation waits on the last one's store. The PR's loop keepsnext_loc_in a register and only stores it.new vari: develop spills the vari pointer to the stack every iteration for the vector grow path. That grow path is now in the cold function so the spill is gone.reverse_pass_callback: on gcc the placement new drops the null check thatnewwith anoexceptoperator newhas to do, and the spill is gone like above. On clang, develop reloadsChainableStack::instance_after constructing the vari to push it on the var stack, and the PR reuses the one it already loaded.new vari: the loop is basically the same in both, clang was already inlining the vari constructor.alloc/alloc_array: the instructions are identical in both builds. The difference in this table is from where the loop happened to land in the binary, so I would not count it.gcc's develop numbers also moved a few percent between runs on this shared machine (e.g.
new variwas 2.75 in an earlier run), so take the gcc vari / callback numbers as roughly 5-15%.I have a PR that turns off our auto performance testing models, but I will be pretty interested in whether they show a noticeable speedup here.
Running the benchmarks
The benchmark code and instructions for downloading and running it are in this gist. It builds the benchmark against this branch and its merge base with develop and compares them with Google Benchmark's
compare.py.Tests
New tests leave the arena 8 bytes past an aligned boundary and then capture a
Vector2d, aMatrix4d, and fixed-sizevarmatrices inreverse_pass_callback,make_callback_vari, andmultiply. They all segfault on develop.Release notes
Fixed a crash when a reverse mode callback captures types that need more than 8 byte alignment. The arena allocator now supports aligned allocations and is faster with clang.
Checklist
Copyright holder: Simons Foundation
The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
- Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
- Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)
the basic tests are passing
./runTests.py test/unit)make test-headers)make test-math-dependencies)make doxygen)make cpplint)the code is written in idiomatic C++ and changes are documented in the doxygen
the new changes are tested