From 6cf558114a2af1cb7d51b9c2c1b6dfbd4fb4bc3a Mon Sep 17 00:00:00 2001 From: Nick Thompson Date: Mon, 2 Nov 2026 19:28:17 -0800 Subject: [PATCH] Stop pre-constructing reverse-mode tape blocks, and reuse them after a rewind Blocks are raw storage now, so a second-order tape adds no leaves to the first-order tape, and a rewound tape refills its blocks instead of allocating new ones. Documents that rvar copies share a node. Closes #1624 Co-Authored-By: Claude Opus 5.5 --- doc/differentiation/autodiff_reverse.qbk | 2 + ...everse_mode_autodiff_memory_management.hpp | 124 ++++++------------ ...se_mode_autodiff_flat_linear_allocator.cpp | 77 +++++++++++ 3 files changed, 116 insertions(+), 87 deletions(-) diff --git a/doc/differentiation/autodiff_reverse.qbk b/doc/differentiation/autodiff_reverse.qbk index d438c0e6fb..ce964640cd 100644 --- a/doc/differentiation/autodiff_reverse.qbk +++ b/doc/differentiation/autodiff_reverse.qbk @@ -174,6 +174,8 @@ Some key points about reverse mode automatic differentiation: 3. While reverse mode is fast for computing the gradient, it has to store all the intermediate values from the forward pass to be used during the backward pass. This can be a significant memory overhead, especially for deep networks or complex functions. +4. A copy of an `rvar` shares its node on the tape, even after one of them is assigned a new number, so derivatives with respect to the two are added together. So `std::vector> x(n, rvar(0))` holds one variable n times. Write `std::vector> x(n)`, or assign each element a new `rvar`, to get n independent variables. + [h1:overloaded-functions List of overloaded functions] [table [[Function] [argument types] [implementation] [return type] [left derivative] [right derivative]] diff --git a/include/boost/math/differentiation/detail/reverse_mode_autodiff_memory_management.hpp b/include/boost/math/differentiation/detail/reverse_mode_autodiff_memory_management.hpp index f562a6f1ab..e5266e80b4 100644 --- a/include/boost/math/differentiation/detail/reverse_mode_autodiff_memory_management.hpp +++ b/include/boost/math/differentiation/detail/reverse_mode_autodiff_memory_management.hpp @@ -214,8 +214,15 @@ class flat_linear_allocator public: // store vector of unique pointers to arrays // to avoid vector reference invalidation - using buffer_type = std::array; - using buffer_ptr = std::unique_ptr>; + // Uninitialized storage: elements are constructed by emplace_back. Default-constructing a whole block + // would, for the rvar elements of a second-order tape, put buffer_size leaves on the first-order tape. + struct buffer_type + { + alignas(RealType) unsigned char bytes[sizeof(RealType) * buffer_size]; + RealType& operator[](size_t i) noexcept { return *reinterpret_cast(bytes + sizeof(RealType) * i); } + const RealType& operator[](size_t i) const noexcept { return *reinterpret_cast(bytes + sizeof(RealType) * i); } + }; + using buffer_ptr = std::unique_ptr; private: std::vector data_; @@ -241,13 +248,25 @@ class flat_linear_allocator size_t item_id() const noexcept { return total_size_ % buffer_size; } private: - void allocate_buffer() + // The storage for element i; blocks are allocated on first use and kept, so a rewound tape reuses them. + RealType* slot(size_t i) { - data_.emplace_back(std::make_unique()); + while (i / buffer_size >= data_.size()) { + data_.emplace_back(buffer_ptr(new buffer_type)); + } + return &(*data_[i / buffer_size])[i % buffer_size]; + } + // Destroys the elements from n on; their blocks are kept. + void truncate(size_t n) + { + for (size_t i = n; i < total_size_; ++i) { + (*this)[i].~RealType(); + } + total_size_ = n; } public: - flat_linear_allocator() { allocate_buffer(); } + flat_linear_allocator() { slot(0); } flat_linear_allocator(const flat_linear_allocator&) = delete; flat_linear_allocator& operator=(const flat_linear_allocator&) = delete; flat_linear_allocator(flat_linear_allocator&&) = delete; @@ -271,15 +290,16 @@ class flat_linear_allocator */ void clear() { + destroy_all(); data_.clear(); total_size_ = 0; checkpoints_.clear(); - allocate_buffer(); + slot(0); } // doesn't delete anything, only sets the current index to zero - void reset() { total_size_ = 0; } - void rewind() { total_size_ = 0; }; + void reset() { truncate(0); } + void rewind() { truncate(0); }; // adds current index as a checkpoint to be able to walk back to void add_checkpoint() @@ -297,10 +317,10 @@ class flat_linear_allocator * */ void reset_checkpoints() { checkpoints_.clear(); } - void rewind_to_last_checkpoint() { total_size_ = checkpoints_.back(); } + void rewind_to_last_checkpoint() { truncate(checkpoints_.back()); } void rewind_to_checkpoint_at(size_t index) { - total_size_ = checkpoints_[index]; + truncate(checkpoints_[index]); } void fill(const RealType& val) @@ -312,34 +332,12 @@ class flat_linear_allocator } } - /** @brief emplaces back object at the end of the - * data structure, calls default constructor */ - iterator emplace_back() - { - if (item_id() == 0 && total_size_ != 0) { - allocate_buffer(); - } - size_t bid = buffer_id(); - size_t iid = item_id(); - - RealType* ptr = &(*data_[bid])[iid]; - new (ptr) RealType(); - ++total_size_; - return iterator(this, total_size_ - 1); - }; - - /** @brief, emplaces back object at end of data structure, + /** @brief emplaces back object at end of data structure, * passes arguments to constructor */ template iterator emplace_back(Args&&... args) { - if (item_id() == 0 && total_size_ != 0) { - allocate_buffer(); - } - BOOST_MATH_ASSERT(buffer_id() < data_.size()); - BOOST_MATH_ASSERT(item_id() < buffer_size); - RealType* ptr = &(*data_[buffer_id()])[item_id()]; - new (ptr) RealType(std::forward(args)...); + new (slot(total_size_)) RealType(std::forward(args)...); ++total_size_; return iterator(this, total_size_ - 1); } @@ -348,66 +346,18 @@ class flat_linear_allocator template iterator emplace_back_n() { - size_t bid = buffer_id(); - size_t iid = item_id(); - if (iid + n < buffer_size) { - RealType* ptr = &(*data_[bid])[iid]; - for (size_t i = 0; i < n; ++i) { - new (ptr + i) RealType(); - } - total_size_ += n; - return iterator(this, total_size_ - n, total_size_ - n, total_size_); - } else { - size_t allocs_in_curr_buffer = buffer_size - iid; - size_t allocs_in_next_buffer = n - (buffer_size - iid); - RealType* ptr = &(*data_[bid])[iid]; - for (size_t i = 0; i < allocs_in_curr_buffer; ++i) { - new (ptr + i) RealType(); - } - allocate_buffer(); - bid = data_.size() - 1; - iid = 0; - total_size_ += n; - - RealType* ptr2 = &(*data_[bid])[iid]; - for (size_t i = 0; i < allocs_in_next_buffer; i++) { - new (ptr2 + i) RealType(); - } - return iterator(this, total_size_ - n, total_size_ - n, total_size_); - } + return emplace_back_n(n); } /** @brief default constructs n objects at end of * data structure, n known at run time */ iterator emplace_back_n(size_t n) { - size_t bid = buffer_id(); - size_t iid = item_id(); - if (iid + n < buffer_size) { - RealType* ptr = &(*data_[bid])[iid]; - for (size_t i = 0; i < n; ++i) { - new (ptr + i) RealType(); - } - total_size_ += n; - return iterator(this, total_size_ - n, total_size_ - n, total_size_); - } else { - size_t allocs_in_curr_buffer = buffer_size - iid; - size_t allocs_in_next_buffer = n - (buffer_size - iid); - RealType* ptr = &(*data_[bid])[iid]; - for (size_t i = 0; i < allocs_in_curr_buffer; ++i) { - new (ptr + i) RealType(); - } - allocate_buffer(); - bid = data_.size() - 1; - iid = 0; - total_size_ += n; - - RealType* ptr2 = &(*data_[bid])[iid]; - for (size_t i = 0; i < allocs_in_next_buffer; i++) { - new (ptr2 + i) RealType(); - } - return iterator(this, total_size_ - n, total_size_ - n, total_size_); + size_t first = total_size_; + for (size_t i = 0; i < n; ++i) { + emplace_back(); } + return iterator(this, first, first, total_size_); } /** @brief number of elements */ diff --git a/test/test_reverse_mode_autodiff_flat_linear_allocator.cpp b/test/test_reverse_mode_autodiff_flat_linear_allocator.cpp index 7e68ad2e4a..18a270a3d1 100644 --- a/test/test_reverse_mode_autodiff_flat_linear_allocator.cpp +++ b/test/test_reverse_mode_autodiff_flat_linear_allocator.cpp @@ -128,10 +128,87 @@ void flat_linear_allocator_test_checkpointing() CHECK_EQUAL(float_allocator.capacity(), 2 * buffer_size); } +template +void flat_linear_allocator_test_block_reuse() +{ + constexpr size_t buffer_size = 16; + rdiff_detail::flat_linear_allocator a{}; + // A rewound allocator refills its blocks instead of allocating new ones. + for (int cycle = 0; cycle < 4; ++cycle) { + if (cycle % 2 == 0) { + for (size_t i = 0; i < 2 * buffer_size + 1; ++i) { + a.emplace_back(T(i)); + } + } else { + auto it = a.emplace_back_n(2 * buffer_size + 1); + for (size_t i = 0; i < 2 * buffer_size + 1; ++i) { + *(it + i) = T(i); + } + } + CHECK_EQUAL(a.capacity(), 3 * buffer_size); + for (size_t i = 0; i < a.size(); ++i) { + CHECK_EQUAL(a[i], T(i)); + } + a.rewind(); + } + // emplace_back_n ending on a block boundary allocates no extra block, and can span several blocks. + a.clear(); + a.emplace_back_n(buffer_size); + CHECK_EQUAL(a.capacity(), buffer_size); + a.emplace_back(T(-1)); + auto it = a.emplace_back_n(2 * buffer_size + 3); + for (size_t i = 0; i < 2 * buffer_size + 3; ++i) { + *(it + i) = T(i); + } + for (size_t i = 0; i < 2 * buffer_size + 3; ++i) { + CHECK_EQUAL(a[buffer_size + 1 + i], T(i)); + } + CHECK_EQUAL(a.capacity(), 4 * buffer_size); +} + +struct counted +{ + static int live; + counted() { ++live; } + counted(const counted&) = delete; + ~counted() { --live; } +}; +int counted::live = 0; + +// Rewinding destroys the discarded elements, and destruction the rest. +void flat_linear_allocator_test_destruction() +{ + { + rdiff_detail::flat_linear_allocator a{}; + a.emplace_back_n(20); + a.add_checkpoint(); + a.emplace_back_n(5); + a.rewind_to_last_checkpoint(); + CHECK_EQUAL(counted::live, 20); + a.emplace_back_n(7); + a.rewind(); + CHECK_EQUAL(counted::live, 0); + a.emplace_back_n(3); + } + CHECK_EQUAL(counted::live, 0); +} + +template +void second_order_tape_adds_no_first_order_nodes() +{ + auto& tape1 = rdiff::get_active_tape(); + auto before = tape1.end() - tape1.begin(); + (void)rdiff::get_active_tape(); + CHECK_EQUAL(tape1.end() - tape1.begin(), before); +} + int main() { boost::mp11::mp_for_each>([](auto t) { flat_linear_allocator_constructors(); }); boost::mp11::mp_for_each>([](auto t) { flat_linear_allocator_test_emplace(); }); boost::mp11::mp_for_each>([](auto t) { flat_linear_allocator_test_checkpointing(); }); + boost::mp11::mp_for_each>>([](auto t) { flat_linear_allocator_test_block_reuse(); }); + flat_linear_allocator_test_destruction(); + boost::mp11::mp_for_each>>([](auto t) { second_order_tape_adds_no_first_order_nodes(); }); return boost::math::test::report_errors(); }