diff --git a/stan/math/memory/stack_alloc.hpp b/stan/math/memory/stack_alloc.hpp index 68227232f10..85f4778d339 100644 --- a/stan/math/memory/stack_alloc.hpp +++ b/stan/math/memory/stack_alloc.hpp @@ -26,7 +26,7 @@ namespace math { * @tparam Type of object to which pointer points. */ template -inline bool is_aligned(T* ptr, unsigned int bytes_aligned) { +STAN_STRONG_INLINE bool is_aligned(T* ptr, unsigned int bytes_aligned) { return (reinterpret_cast(ptr) % bytes_aligned) == 0U; } @@ -35,18 +35,32 @@ const size_t DEFAULT_INITIAL_NBYTES = 1 << 16; // 64KB // FIXME: enforce alignment // big fun to inline, but only called twice -inline char* eight_byte_aligned_malloc(size_t size) { +STAN_STRONG_INLINE char* eight_byte_aligned_malloc(size_t size) { char* ptr = static_cast(malloc(size)); - if (!ptr) { + if (unlikely(!ptr)) { return ptr; // malloc failed to alloc - } - if (!is_aligned(ptr, 8U)) { + } else if (unlikely(!is_aligned(ptr, 8U))) { std::stringstream s; s << "invalid alignment to 8 bytes, ptr=" << reinterpret_cast(ptr) << std::endl; throw std::runtime_error(s.str()); + } else { + return ptr; } - return ptr; +} + +/** + * Return the first address at or after `ptr` that is aligned to + * `alignment` bytes. + * + * @param ptr Pointer to align. + * @param alignment Alignment in bytes, a power of two. + */ +STAN_STRONG_INLINE char* align_up(char* ptr, size_t alignment) { + const uintptr_t address = reinterpret_cast(ptr); + const uintptr_t aligned + = (address + alignment - 1) & ~static_cast(alignment - 1); + return ptr + (aligned - address); } } // namespace internal @@ -64,10 +78,8 @@ inline char* eight_byte_aligned_malloc(size_t size) { * recovered, with the blocks being reused, or all blocks may be * freed, resetting the stack of blocks to its original state. * - * Alignment up to 8 byte boundaries guaranteed for the first malloc, - * and after that it's up to the caller. On 64-bit architectures, - * all struct values should be padded to 8-byte boundaries if they - * contain an 8-byte member or a virtual function. + * Every allocation is at least 8 byte aligned. Larger alignments can be + * requested through the `alignment` argument of `alloc()`. */ class stack_alloc { private: @@ -86,24 +98,26 @@ class stack_alloc { /** * Moves us to the next block of memory, allocating that block * if necessary, and allocates len bytes of memory within that - * block. + * block, aligned to `alignment` bytes. * * @param len Number of bytes to allocate. + * @param alignment Alignment in bytes, a power of two. * @return A pointer to the allocated memory. */ - char* move_to_next_block(size_t len) { - char* result; + STAN_STRONG_INLINE char* move_to_next_block(size_t len, size_t alignment) { + // Blocks are 8 byte aligned, so leave room to align the start. + const size_t needed = alignment > 8 ? len + alignment - 8 : len; ++cur_block_; - // Find the next block (if any) containing at least len bytes. - while ((cur_block_ < blocks_.size()) && (sizes_[cur_block_] < len)) { + // Find the next block (if any) containing at least needed bytes. + while ((cur_block_ < blocks_.size()) && (sizes_[cur_block_] < needed)) { ++cur_block_; } // Allocate a new block if necessary. if (unlikely(cur_block_ >= blocks_.size())) { - // New block should be max(2*size of last block, len) bytes. + // New block should be max(2*size of last block, needed) bytes. size_t newsize = sizes_.back() * 2; - if (newsize < len) { - newsize = len; + if (newsize < needed) { + newsize = needed; } blocks_.push_back(internal::eight_byte_aligned_malloc(newsize)); if (!blocks_.back()) { @@ -111,10 +125,10 @@ class stack_alloc { } sizes_.push_back(newsize); } - result = blocks_[cur_block_]; + char* result = internal::align_up(blocks_[cur_block_], alignment); // Get the object's state back in order. next_loc_ = result + len; - cur_block_end_ = result + sizes_[cur_block_]; + cur_block_end_ = blocks_[cur_block_] + sizes_[cur_block_]; return result; } @@ -158,25 +172,27 @@ class stack_alloc { * Return a newly allocated block of memory of the appropriate * size managed by the stack allocator. * - * The allocated pointer will be 8-byte aligned. If the number - * of bytes requested is not a multiple of 8, the reserved space - * will be padded up to the next multiple of 8. + * The allocated pointer will be aligned to `alignment` bytes, and at + * least 8 bytes. If the number of bytes requested is not a multiple of + * 8, the reserved space will be padded up to the next multiple of 8. * * This function may call C++'s malloc() function, * with any exceptions percolated through this function. * * @param len Number of bytes to allocate. + * @param alignment Alignment in bytes, a power of two. * @return A pointer to the allocated memory. */ - inline void* alloc(size_t len) { + STAN_STRONG_INLINE void* alloc(size_t len, size_t alignment = 8) { size_t pad = len % 8 == 0 ? 0 : 8 - len % 8; // Typically, just return and increment the next location. - char* result = next_loc_; - next_loc_ += len + pad; + char* result + = alignment <= 8 ? next_loc_ : internal::align_up(next_loc_, alignment); + next_loc_ = result + len + pad; // Occasionally, we have to switch blocks. if (unlikely(next_loc_ >= cur_block_end_)) { - result = move_to_next_block(len + pad); + result = move_to_next_block(len + pad, alignment); } return reinterpret_cast(result); } @@ -190,7 +206,7 @@ class stack_alloc { * @return new array allocated on the arena. */ template - inline T* alloc_array(size_t n) { + STAN_STRONG_INLINE T* alloc_array(size_t n) { return static_cast(alloc(n * sizeof(T))); } diff --git a/stan/math/prim/meta/compiler_attributes.hpp b/stan/math/prim/meta/compiler_attributes.hpp index 0381b59bdd3..ab584316d76 100644 --- a/stan/math/prim/meta/compiler_attributes.hpp +++ b/stan/math/prim/meta/compiler_attributes.hpp @@ -2,13 +2,17 @@ #define STAN_MATH_PRIM_META_COMPILER_ATTRIBUTES_HPP #ifdef __GNUC__ + #ifndef likely #define likely(x) __builtin_expect(!!(x), 1) #endif + #ifndef unlikely #define unlikely(x) __builtin_expect(!!(x), 0) #endif + #ifdef __has_attribute + #if __has_attribute(noinline) && __has_attribute(cold) #ifndef STAN_COLD_PATH /** @@ -19,7 +23,30 @@ #define STAN_COLD_PATH __attribute__((noinline, cold)) #endif #endif + +#if __has_attribute(always_inline) +#ifndef STAN_STRONG_INLINE +/** + * Functions tagged with this attribute are always inlined. Use it on small + * hot path functions where the compiler's own inlining heuristics are not + * reliable, such as the arena allocator's fast path. Includes `inline`. + */ +#define STAN_STRONG_INLINE __attribute__((always_inline)) inline +#endif +#endif + +#if __has_attribute(hot) +#ifndef STAN_HOT_PATH +/** + * Tells compiler this is a hot path function that should be more + * aggressively optimized. + */ +#define STAN_HOT_PATH __attribute__((hot)) +#endif +#endif + #endif + #ifndef STAN_DEPRECATED #if __has_attribute(deprecated) #define STAN_DEPRECATED __attribute__((deprecated)) @@ -30,15 +57,26 @@ #endif -#ifndef STAN_COLD_PATH -#define STAN_COLD_PATH -#endif #ifndef likely #define likely(x) x #endif + #ifndef unlikely #define unlikely(x) x #endif + +#ifndef STAN_COLD_PATH +#define STAN_COLD_PATH +#endif + +#ifndef STAN_STRONG_INLINE +#define STAN_STRONG_INLINE inline +#endif + +#ifndef STAN_HOT_PATH +#define STAN_HOT_PATH +#endif + #ifndef STAN_DEPRECATED #define STAN_DEPRECATED #endif diff --git a/stan/math/rev/core/autodiffstackstorage.hpp b/stan/math/rev/core/autodiffstackstorage.hpp index fecdab186b5..21bcea595d7 100644 --- a/stan/math/rev/core/autodiffstackstorage.hpp +++ b/stan/math/rev/core/autodiffstackstorage.hpp @@ -110,6 +110,36 @@ struct AutodiffStackSingleton { std::vector nested_var_stack_sizes_; std::vector nested_var_nochain_stack_sizes_; std::vector nested_var_alloc_stack_starts_; + + STAN_STRONG_INLINE void push_var_stack(ChainableT *vi) { + push(var_stack_, vi); + } + STAN_STRONG_INLINE void push_var_nochain_stack(ChainableT *vi) { + push(var_nochain_stack_, vi); + } + + /** + * Once we create memory in `stack` we do not release it till + * the end of the program. + * gcc/clang's `vector::push_back` does the same check as below, but here + * we force inlining. gcc inlines `push_back` on both sides of our if, + * sees we do the capacity check it does, and then removes its capacity + * check. So then we get our nice `if(likely(...))`. + * This optimization depends on `stack.push_back` being inlined. + * Hence why the function is strongly inlined and hot. + */ + STAN_STRONG_INLINE STAN_HOT_PATH static void push( + std::vector &stack, ChainableT *vi) { + if (likely(stack.size() != stack.capacity())) { + stack.push_back(vi); + } else { + grow_and_push(stack, vi); + } + } + STAN_COLD_PATH static void grow_and_push(std::vector &stack, + ChainableT *vi) { + stack.push_back(vi); + } }; explicit AutodiffStackSingleton(AutodiffStackSingleton_t const &) = delete; diff --git a/stan/math/rev/core/callback_vari.hpp b/stan/math/rev/core/callback_vari.hpp index 6ecd88da8e2..a33fa60c1bc 100644 --- a/stan/math/rev/core/callback_vari.hpp +++ b/stan/math/rev/core/callback_vari.hpp @@ -2,6 +2,7 @@ #define STAN_MATH_REV_CORE_CALLBACK_VARI_HPP #include +#include #include namespace stan { @@ -40,8 +41,12 @@ struct callback_vari : public vari_value { template inline internal::callback_vari, F>* make_callback_vari( T&& value, F&& functor) { - return new internal::callback_vari, F>( - std::move(value), std::forward(functor)); + // Captures such as fixed-size Eigen matrices can need more than the 8 byte + // alignment of `vari_base::operator new`. + using vari_t = internal::callback_vari, F>; + return ::new (ChainableStack::instance_->memalloc_.alloc(sizeof(vari_t), + alignof(vari_t))) + vari_t(std::move(value), std::forward(functor)); } /** diff --git a/stan/math/rev/core/reverse_pass_callback.hpp b/stan/math/rev/core/reverse_pass_callback.hpp index b93ede8c3c1..c233b805ef0 100644 --- a/stan/math/rev/core/reverse_pass_callback.hpp +++ b/stan/math/rev/core/reverse_pass_callback.hpp @@ -2,6 +2,7 @@ #define STAN_MATH_REV_FUNCTOR_REVERSE_PASS_CALLBACK_HPP #include +#include namespace stan { namespace math { @@ -11,9 +12,10 @@ template struct reverse_pass_callback_vari : public vari_base { F rev_functor_; - explicit reverse_pass_callback_vari(F&& rev_functor) + STAN_STRONG_INLINE STAN_HOT_PATH explicit reverse_pass_callback_vari( + F&& rev_functor) : rev_functor_(std::forward(rev_functor)) { - ChainableStack::instance_->var_stack_.push_back(this); + ChainableStack::instance_->push_var_stack(this); } inline void chain() final { rev_functor_(); } @@ -35,8 +37,12 @@ struct reverse_pass_callback_vari : public vari_base { * @param functor funtor or other callable to call in the reverse pass */ template -inline void reverse_pass_callback(F&& functor) { - new internal::reverse_pass_callback_vari(std::forward(functor)); +STAN_STRONG_INLINE STAN_HOT_PATH void reverse_pass_callback(F&& functor) { + // Captures such as fixed-size Eigen matrices can need more than the 8 byte + // alignment of `vari_base::operator new`. + using vari_t = internal::reverse_pass_callback_vari; + ::new (ChainableStack::instance_->memalloc_.alloc( + sizeof(vari_t), alignof(vari_t))) vari_t(std::forward(functor)); } } // namespace math diff --git a/stan/math/rev/core/vari.hpp b/stan/math/rev/core/vari.hpp index 0e9f6047ba8..c6977e353d6 100644 --- a/stan/math/rev/core/vari.hpp +++ b/stan/math/rev/core/vari.hpp @@ -44,7 +44,7 @@ class vari_base { * @param nbytes Number of bytes to allocate. * @return Pointer to allocated bytes. */ - static inline void* operator new(size_t nbytes) noexcept { + STAN_STRONG_INLINE static void* operator new(size_t nbytes) noexcept { return ChainableStack::instance_->memalloc_.alloc(nbytes); } @@ -103,8 +103,8 @@ class vari_value>> : public vari_base { * @param x Value of the constructed variable. */ template * = nullptr> - vari_value(S x) noexcept : val_(x) { // NOLINT - ChainableStack::instance_->var_stack_.push_back(this); + STAN_STRONG_INLINE vari_value(S x) noexcept : val_(x) { // NOLINT + ChainableStack::instance_->push_var_stack(this); } /** @@ -123,11 +123,11 @@ class vari_value>> : public vari_base { * that its `chain()` method is not called. */ template * = nullptr> - vari_value(S x, bool stacked) noexcept : val_(x) { + STAN_STRONG_INLINE vari_value(S x, bool stacked) noexcept : val_(x) { if (stacked) { - ChainableStack::instance_->var_stack_.push_back(this); + ChainableStack::instance_->push_var_stack(this); } else { - ChainableStack::instance_->var_nochain_stack_.push_back(this); + ChainableStack::instance_->push_var_nochain_stack(this); } } diff --git a/test/unit/math/memory/stack_alloc_test.cpp b/test/unit/math/memory/stack_alloc_test.cpp index b395df387e9..7764c36a0fd 100644 --- a/test/unit/math/memory/stack_alloc_test.cpp +++ b/test/unit/math/memory/stack_alloc_test.cpp @@ -100,6 +100,43 @@ TEST(stack_alloc, alloc_aligned) { allocator.recover_all(); } +TEST(stack_alloc, alloc_over_aligned) { + stan::math::stack_alloc allocator; + // leave the next free byte 8 but not 64 byte aligned + allocator.alloc(8); + for (size_t alignment : {16U, 32U, 64U}) { + double* x + = static_cast(allocator.alloc(4 * sizeof(double), alignment)); + EXPECT_TRUE(stan::math::is_aligned(x, alignment)) << alignment; + for (int i = 0; i < 4; ++i) { + x[i] = i; + } + EXPECT_FLOAT_EQ(3.0, x[3]); + // plain allocations after an aligned one stay 8 byte aligned + EXPECT_TRUE(stan::math::is_aligned(allocator.alloc(sizeof(double)), 8U)); + } +} + +TEST(stack_alloc, alloc_over_aligned_new_blocks) { + stan::math::stack_alloc allocator; + // larger than the initial block, so every allocation needs a new block + const size_t big = 1 << 17; + for (int i = 0; i < 4; ++i) { + allocator.alloc(8); + char* x = static_cast(allocator.alloc(big, 64)); + EXPECT_TRUE(stan::math::is_aligned(x, 64U)); + x[0] = 1; + x[big - 1] = 2; + EXPECT_TRUE(allocator.in_stack(x)); + } + // reused blocks after recovering memory are aligned too + allocator.recover_all(); + for (int i = 0; i < 4; ++i) { + allocator.alloc(8); + EXPECT_TRUE(stan::math::is_aligned(allocator.alloc(big, 64), 64U)); + } +} + TEST(stack_alloc, in_stack) { stan::math::stack_alloc allocator; diff --git a/test/unit/math/rev/core/callback_vari_test.cpp b/test/unit/math/rev/core/callback_vari_test.cpp index eb9b4a7740a..a0cc0e7ae2f 100644 --- a/test/unit/math/rev/core/callback_vari_test.cpp +++ b/test/unit/math/rev/core/callback_vari_test.cpp @@ -45,6 +45,25 @@ TEST_F(AgradRev, RevCBack_callback_vari_eigen_test) { EXPECT_MATRIX_EQ(a.adj(), Eigen::MatrixXd::Constant(2, 3, 2)); } +TEST_F(AgradRev, RevCBack_callback_vari_over_aligned_capture) { + // a callback stores its captures inside the vari allocated on the arena + const Eigen::Vector2d m(1.0, 2.0); + stan::math::var a = 1; + bool aligned = false; + // leave the next free byte of the arena 8 bytes past a 16 byte boundary + auto& arena = stan::math::ChainableStack::instance_->memalloc_; + while (!stan::math::is_aligned(arena.alloc(8), 16U)) { + } + stan::math::var b = stan::math::make_callback_vari( + a.val(), [a, m, &aligned](const auto& vi) mutable { + aligned = stan::math::is_aligned(m.data(), alignof(Eigen::Vector2d)); + a.adj() += vi.adj_ * m.sum(); + }); + b.grad(); + EXPECT_TRUE(aligned); + EXPECT_FLOAT_EQ(a.adj(), 3.0); +} + TEST_F(AgradRev, RevCBack_make_callback_var_scalar_test) { stan::math::var a = 1; stan::math::var b = 1; diff --git a/test/unit/math/rev/core/reverse_pass_callback_test.cpp b/test/unit/math/rev/core/reverse_pass_callback_test.cpp index a91317d7f2c..fecf002fc54 100644 --- a/test/unit/math/rev/core/reverse_pass_callback_test.cpp +++ b/test/unit/math/rev/core/reverse_pass_callback_test.cpp @@ -1,7 +1,19 @@ #include #include +#include #include +namespace { +/** + * Leave the next free byte of the arena 8 bytes past a 16 byte boundary. + */ +void misalign_arena() { + auto& arena = stan::math::ChainableStack::instance_->memalloc_; + while (!stan::math::is_aligned(arena.alloc(8), 16U)) { + } +} +} // namespace + TEST_F(AgradRev, Rev_reverse_pass_callback_test) { stan::math::var a = 1; stan::math::var b = 1; @@ -12,3 +24,35 @@ TEST_F(AgradRev, Rev_reverse_pass_callback_test) { b.grad(); EXPECT_EQ(a.adj(), 3); } + +TEST_F(AgradRev, Rev_reverse_pass_callback_over_aligned_capture) { + // a callback stores its captures inside the vari allocated on the arena + const Eigen::Vector2d m(1.0, 2.0); + stan::math::var b = 1; + bool captured_aligned = false; + double captured_sum = 0; + misalign_arena(); + stan::math::reverse_pass_callback([m, &captured_aligned, &captured_sum]() { + captured_aligned + = stan::math::is_aligned(m.data(), alignof(Eigen::Vector2d)); + captured_sum = m.sum(); + }); + b.grad(); + EXPECT_TRUE(captured_aligned); + EXPECT_FLOAT_EQ(captured_sum, 3.0); +} + +TEST_F(AgradRev, Rev_multiply_fixed_size_var_matrices) { + using stan::math::var; + Eigen::Matrix f; + f << 1, 2, 3, 4; + misalign_arena(); + Eigen::Matrix c = stan::math::multiply(f, f); + stan::math::sum(c).grad(); + const Eigen::Matrix2d f_val = f.val(); + EXPECT_MATRIX_NEAR(c.val(), f_val * f_val, 1e-12); + // d sum(F * F) / dF = ones * F^T + F^T * ones + const Eigen::Matrix2d ones = Eigen::Matrix2d::Ones(); + EXPECT_MATRIX_NEAR( + f.adj(), ones * f_val.transpose() + f_val.transpose() * ones, 1e-12); +}