diff --git a/stan/math/memory/stack_alloc.hpp b/stan/math/memory/stack_alloc.hpp index 68227232f10..261e08c1563 100644 --- a/stan/math/memory/stack_alloc.hpp +++ b/stan/math/memory/stack_alloc.hpp @@ -10,6 +10,29 @@ #include #include #include +#ifdef _WIN32 +#include +#endif + +/** + * `STAN_ALIGNED_ALLOC(alignment, size)` allocates `size` bytes aligned to + * `alignment` bytes and `STAN_ALIGNED_FREE(ptr)` releases that memory. + * + * MSVC and MinGW do not provide `std::aligned_alloc`, so on Windows these + * use `_aligned_malloc` and `_aligned_free`. Memory from `_aligned_malloc` + * must not be released with `free()`. Defining `STAN_ALIGNED_ALLOC` before + * including Stan replaces both, so `STAN_ALIGNED_FREE` must be defined too. + */ +#ifndef STAN_ALIGNED_ALLOC +#ifdef _WIN32 +#define STAN_ALIGNED_ALLOC(alignment, size) _aligned_malloc((size), (alignment)) +#define STAN_ALIGNED_FREE(ptr) _aligned_free(ptr) +#else +#define STAN_ALIGNED_ALLOC(alignment, size) \ + std::aligned_alloc((alignment), (size)) +#define STAN_ALIGNED_FREE(ptr) std::free(ptr) +#endif +#endif namespace stan { namespace math { @@ -33,20 +56,33 @@ inline bool is_aligned(T* ptr, unsigned int bytes_aligned) { namespace internal { 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) { - char* ptr = static_cast(malloc(size)); - if (!ptr) { - return ptr; // malloc failed to alloc - } - if (!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()); - } - return ptr; +/** + * Allocate a block of memory aligned to `alignment` bytes with + * `STAN_ALIGNED_ALLOC`. The block must be released with + * `STAN_ALIGNED_FREE`. + * + * @param size Number of bytes to allocate; rounded up to a multiple of + * `alignment` as `std::aligned_alloc` requires. + * @param alignment Alignment in bytes, a power of two. + * @return Pointer to the block, or `nullptr` if the allocation failed. + */ +inline char* aligned_malloc(size_t size, size_t alignment) { + const size_t rounded_size = (size + alignment - 1) / alignment * alignment; + return static_cast(STAN_ALIGNED_ALLOC(alignment, rounded_size)); +} + +/** + * 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. + */ +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 +100,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: @@ -83,19 +117,32 @@ class stack_alloc { std::vector nested_next_locs_; std::vector nested_cur_block_ends_; + /** + * Return `true` if block `i` has room for `len` bytes starting at an + * address aligned to `alignment` bytes. + */ + inline bool block_fits(size_t i, size_t len, size_t alignment) const { + const size_t offset + = internal::align_up(blocks_[i], alignment) - blocks_[i]; + return offset + len <= sizes_[i]; + } + /** * 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. New blocks come + * from `std::aligned_alloc` with this alignment. * @return A pointer to the allocated memory. */ - char* move_to_next_block(size_t len) { - char* result; + STAN_STRONG_INLINE inline char* move_to_next_block(size_t len, + size_t alignment = 8) { ++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) with room for len aligned bytes. + while ((cur_block_ < blocks_.size()) + && !block_fits(cur_block_, len, alignment)) { ++cur_block_; } // Allocate a new block if necessary. @@ -105,16 +152,16 @@ class stack_alloc { if (newsize < len) { newsize = len; } - blocks_.push_back(internal::eight_byte_aligned_malloc(newsize)); + blocks_.push_back(internal::aligned_malloc(newsize, alignment)); if (!blocks_.back()) { throw std::bad_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; } @@ -125,11 +172,10 @@ class stack_alloc { * * @param initial_nbytes Initial number of bytes for the * allocator. Defaults to (1 << 16) = 64KB initial bytes. - * @throws std::runtime_error if the underlying malloc is not 8-byte - * aligned. + * @throws std::bad_alloc if the initial block cannot be allocated. */ explicit stack_alloc(size_t initial_nbytes = internal::DEFAULT_INITIAL_NBYTES) - : blocks_(1, internal::eight_byte_aligned_malloc(initial_nbytes)), + : blocks_(1, internal::aligned_malloc(initial_nbytes, 8)), sizes_(1, initial_nbytes), cur_block_(0), cur_block_end_(blocks_[0] + initial_nbytes), @@ -149,36 +195,39 @@ class stack_alloc { // free ALL blocks for (auto& block : blocks_) { if (block) { - free(block); + STAN_ALIGNED_FREE(block); } } } /** - * Return a newly allocated block of memory of the appropriate - * size managed by the stack allocator. + * Return a newly allocated block of memory of the appropriate size + * managed by the stack allocator, aligned to `alignment` bytes. * - * 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. + * Every allocation is at least 8 byte aligned; an `alignment` below 8 + * gives 8 byte alignment. The reserved space is padded up to a multiple + * of 8 bytes, so allocations after this one stay 8 byte aligned. * * This function may call C++'s malloc() function, * with any exceptions percolated through this function. * - * @param len Number of bytes to allocate. + * @param[in] len Number of bytes to allocate. + * @param[in] alignment Alignment in bytes, a power of two. * @return A pointer to the allocated memory. */ - inline void* alloc(size_t len) { - size_t pad = len % 8 == 0 ? 0 : 8 - len % 8; - + STAN_STRONG_INLINE inline void* alloc(size_t len, size_t alignment = 8) { + const size_t padded_len = (len + 7) & ~static_cast(7); + // Every block and every allocation is 8 byte aligned, so only + // larger alignments need to round up the next location. + char* result + = alignment <= 8 ? next_loc_ : internal::align_up(next_loc_, alignment); // Typically, just return and increment the next location. - char* result = next_loc_; - next_loc_ += len + pad; + next_loc_ = result + padded_len; // 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(padded_len, alignment); } - return reinterpret_cast(result); + return result; } /** @@ -187,11 +236,12 @@ class stack_alloc { * * @tparam T type of entries in allocated array. * @param[in] n size of array to allocate. + * @param[in] alignment Alignment in bytes, a power of two. * @return new array allocated on the arena. */ template - inline T* alloc_array(size_t n) { - return static_cast(alloc(n * sizeof(T))); + STAN_STRONG_INLINE inline T* alloc_array(size_t n, size_t alignment = 8) { + return static_cast(alloc(n * sizeof(T), alignment)); } /** @@ -243,7 +293,7 @@ class stack_alloc { // frees all BUT the first (index 0) block for (size_t i = 1; i < blocks_.size(); ++i) { if (blocks_[i]) { - free(blocks_[i]); + STAN_ALIGNED_FREE(blocks_[i]); } } sizes_.resize(1); diff --git a/stan/math/opencl/rev/vari.hpp b/stan/math/opencl/rev/vari.hpp index d51d1454cf3..d0ee5349791 100644 --- a/stan/math/opencl/rev/vari.hpp +++ b/stan/math/opencl/rev/vari.hpp @@ -207,7 +207,7 @@ class vari_value> : public chainable_alloc, explicit vari_value(S&& x) : chainable_alloc(), vari_cl_base(std::forward(x), constant(0, x.rows(), x.cols())) { - ChainableStack::instance_->var_stack_.push_back(this); + ChainableStack::instance_->push_var_stack(this); } /** @@ -227,7 +227,7 @@ class vari_value> : public chainable_alloc, require_vt_same* = nullptr> explicit vari_value(const S& x) : chainable_alloc(), vari_cl_base(x, constant(0, x.rows(), x.cols())) { - ChainableStack::instance_->var_nochain_stack_.push_back(this); + ChainableStack::instance_->push_var_nochain_stack(this); } /** @@ -253,9 +253,9 @@ class vari_value> : public chainable_alloc, : chainable_alloc(), vari_cl_base(std::forward(x), constant(0, x.rows(), x.cols())) { 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); } } @@ -276,7 +276,7 @@ class vari_value> : public chainable_alloc, explicit vari_value(S&& val, K&& adj) : chainable_alloc(), vari_cl_base(std::forward(val), std::forward(adj)) { - ChainableStack::instance_->var_nochain_stack_.push_back(this); + ChainableStack::instance_->push_var_nochain_stack(this); } /** diff --git a/stan/math/prim/meta/compiler_attributes.hpp b/stan/math/prim/meta/compiler_attributes.hpp index 0381b59bdd3..f17d2952727 100644 --- a/stan/math/prim/meta/compiler_attributes.hpp +++ b/stan/math/prim/meta/compiler_attributes.hpp @@ -19,6 +19,16 @@ #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. + */ +#define STAN_STRONG_INLINE __attribute__((always_inline)) +#endif +#endif #endif #ifndef STAN_DEPRECATED #if __has_attribute(deprecated) @@ -33,6 +43,9 @@ #ifndef STAN_COLD_PATH #define STAN_COLD_PATH #endif +#ifndef STAN_STRONG_INLINE +#define STAN_STRONG_INLINE +#endif #ifndef likely #define likely(x) x #endif diff --git a/stan/math/rev/core/autodiffstackstorage.hpp b/stan/math/rev/core/autodiffstackstorage.hpp index fecdab186b5..a998c7d0935 100644 --- a/stan/math/rev/core/autodiffstackstorage.hpp +++ b/stan/math/rev/core/autodiffstackstorage.hpp @@ -110,6 +110,44 @@ struct AutodiffStackSingleton { std::vector nested_var_stack_sizes_; std::vector nested_var_nochain_stack_sizes_; std::vector nested_var_alloc_stack_starts_; + + /** + * Push a vari onto `var_stack_`. + * + * The stacks are reused across gradient evaluations, so they almost + * never need to grow. Growing is kept out of line so the inlined push + * is only a compare, a store and an increment. + * + * @param vi vari to push. + */ + STAN_STRONG_INLINE inline void push_var_stack(ChainableT *vi) { + push_back_reused(var_stack_, vi); + } + + /** + * Push a vari onto `var_nochain_stack_`. As for `push_var_stack()`, + * growing the stack is kept out of line. + * + * @param vi vari to push. + */ + STAN_STRONG_INLINE inline void push_var_nochain_stack(ChainableT *vi) { + push_back_reused(var_nochain_stack_, vi); + } + + private: + STAN_STRONG_INLINE static inline void push_back_reused( + std::vector &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 &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..0ff94bd7556 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)); + // The callback can capture types that need more than the 8 byte alignment + // of `vari_base::operator new(size_t)` (see `reverse_pass_callback`). + 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..f9b6b9ab050 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,9 @@ template struct reverse_pass_callback_vari : public vari_base { F rev_functor_; - explicit reverse_pass_callback_vari(F&& rev_functor) + STAN_STRONG_INLINE 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 +36,14 @@ 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 inline void reverse_pass_callback(F&& functor) { + // The callback can capture types that need more than the 8 byte alignment + // of `vari_base::operator new(size_t)`, e.g. fixed-size Eigen matrices. + // C++17 only calls the aligned operator new above + // __STDCPP_DEFAULT_NEW_ALIGNMENT__, so allocate with the type's alignment. + 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..900708bc8f3 100644 --- a/stan/math/rev/core/vari.hpp +++ b/stan/math/rev/core/vari.hpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -44,10 +45,26 @@ 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 inline void* operator new(size_t nbytes) noexcept { return ChainableStack::instance_->memalloc_.alloc(nbytes); } + /** + * Allocate memory from the underlying memory pool for a type that needs + * more alignment than `operator new(size_t)` provides, e.g. a callback + * that captures fixed-size Eigen matrices. C++17 calls this overload for + * such types. + * + * @param nbytes Number of bytes to allocate. + * @param alignment Required alignment in bytes. + * @return Pointer to allocated bytes. + */ + STAN_STRONG_INLINE static inline void* operator new( + size_t nbytes, std::align_val_t alignment) noexcept { + return ChainableStack::instance_->memalloc_.alloc( + nbytes, static_cast(alignment)); + } + /** * Delete a pointer from the underlying memory pool. * @@ -62,6 +79,15 @@ class vari_base { static inline void operator delete( void* /* ignore arg */) noexcept { /* no op */ } + + /** + * Delete a pointer allocated with the aligned `operator new`. As for the + * unaligned version this is a no-op. + */ + static inline void operator delete( + void* /* ignore arg */, std::align_val_t /* ignore arg */) noexcept { + /* no op */ + } }; /** @@ -103,8 +129,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 +149,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); } } @@ -697,7 +723,7 @@ class vari_value, is_eigen_dense_base>> ? x.rows() : x.cols()) { adj_.setZero(); - ChainableStack::instance_->var_nochain_stack_.push_back(this); + ChainableStack::instance_->push_var_nochain_stack(this); } /** @@ -728,9 +754,9 @@ class vari_value, is_eigen_dense_base>> : x.cols()) { adj_.setZero(); 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); } } @@ -749,7 +775,7 @@ class vari_value, is_eigen_dense_base>> template * = nullptr, require_assignable_t* = nullptr> explicit vari_value(const S& val, const K& adj) : val_(val), adj_(adj) { - ChainableStack::instance_->var_nochain_stack_.push_back(this); + ChainableStack::instance_->push_var_nochain_stack(this); } protected: @@ -868,13 +894,13 @@ class vari_value> : public vari_base { val_.innerIndexPtr(), arena_matrix(val_.nonZeros()).setZero().data(), val_.innerNonZeroPtr()) { - ChainableStack::instance_->var_stack_.push_back(this); + ChainableStack::instance_->push_var_stack(this); } vari_value(const arena_matrix& val, const arena_matrix& adj) : val_(val), adj_(adj) { - ChainableStack::instance_->var_stack_.push_back(this); + ChainableStack::instance_->push_var_stack(this); } /** @@ -902,9 +928,9 @@ class vari_value> : public vari_base { arena_matrix(val_.nonZeros()).setZero().data(), val_.innerNonZeroPtr()) { 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/stan/math/rev/functor/cvodes_integrator_adjoint.hpp b/stan/math/rev/functor/cvodes_integrator_adjoint.hpp index d8e030355ee..64e7e10d468 100644 --- a/stan/math/rev/functor/cvodes_integrator_adjoint.hpp +++ b/stan/math/rev/functor/cvodes_integrator_adjoint.hpp @@ -356,7 +356,7 @@ class cvodes_integrator_adjoint_vari : public vari_base { t_init = t_final; } - ChainableStack::instance_->var_stack_.push_back(this); + ChainableStack::instance_->push_var_stack(this); } private: diff --git a/test/unit/math/memory/stack_alloc_test.cpp b/test/unit/math/memory/stack_alloc_test.cpp index b395df387e9..cedefcea306 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; @@ -132,3 +169,17 @@ TEST(stack_alloc, in_stack_second_block) { EXPECT_FALSE(allocator.in_stack(x)); EXPECT_FALSE(allocator.in_stack(y)); } + +TEST(stack_alloc, aligned_alloc_and_free) { + for (std::size_t alignment : {8, 16, 32, 64, 128}) { + void* ptr = STAN_ALIGNED_ALLOC(alignment, 4 * alignment); + ASSERT_NE(ptr, nullptr) << alignment; + EXPECT_TRUE(stan::math::is_aligned(ptr, alignment)) << alignment; + STAN_ALIGNED_FREE(ptr); + } +} + +TEST(stack_alloc, first_block_aligned) { + stan::math::stack_alloc allocator; + EXPECT_TRUE(stan::math::is_aligned(allocator.alloc(1), 8U)); +} diff --git a/test/unit/math/rev/core/callback_vari_test.cpp b/test/unit/math/rev/core/callback_vari_test.cpp index eb9b4a7740a..fede5ee50c7 100644 --- a/test/unit/math/rev/core/callback_vari_test.cpp +++ b/test/unit/math/rev/core/callback_vari_test.cpp @@ -45,6 +45,37 @@ 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 m16(1.0, 2.0); + const Eigen::Matrix4d m32 = Eigen::Matrix4d::Identity(); + stan::math::var a = 1; + bool aligned16 = false; + bool aligned32 = 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, m16, &aligned16](const auto& vi) mutable { + aligned16 + = stan::math::is_aligned(m16.data(), alignof(Eigen::Vector2d)); + a.adj() += vi.adj_ * m16.sum(); + }); + while (!stan::math::is_aligned(arena.alloc(8), 16U)) { + } + stan::math::var c = stan::math::make_callback_vari( + b.val(), [b, m32, &aligned32](const auto& vi) mutable { + aligned32 + = stan::math::is_aligned(m32.data(), alignof(Eigen::Matrix4d)); + b.adj() += vi.adj_ * m32.trace(); + }); + c.grad(); + EXPECT_TRUE(aligned16); + EXPECT_TRUE(aligned32); + EXPECT_FLOAT_EQ(a.adj(), 12.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..c9925ba43cf 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,56 @@ 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) { + // Fixed-size Eigen matrices require more than 8 byte alignment, and a + // callback stores its captures inside the vari allocated on the arena. + const Eigen::Matrix2d m = Eigen::Matrix2d::Identity(); + stan::math::var b = 1; + bool captured_aligned = false; + double captured_sum = 0; + // leave the next free arena byte 8 but not 32 byte aligned + stan::math::ChainableStack::instance_->memalloc_.alloc(8); + stan::math::reverse_pass_callback([m, &captured_aligned, &captured_sum]() { + captured_aligned + = stan::math::is_aligned(m.data(), alignof(Eigen::Matrix2d)); + captured_sum = (m * m).sum(); + }); + b.grad(); + EXPECT_TRUE(captured_aligned); + EXPECT_FLOAT_EQ(captured_sum, 2.0); +} + +TEST_F(AgradRev, Rev_reverse_pass_callback_16_byte_aligned_capture) { + // Vector2d needs 16 byte alignment, which is __STDCPP_DEFAULT_NEW_ALIGNMENT__ + // on common platforms, so C++17 does not call the aligned operator new. + 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; + // leave the next free arena byte 8 but not 32 byte aligned + stan::math::ChainableStack::instance_->memalloc_.alloc(8); + 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); +}