differentiation/detail/reverse_mode_autodiff_memory_management.hpp

0001 //           Copyright Maksym Zhelyenzyakov 2025-2026.
0002 // Distributed under the Boost Software License, Version 1.0.
0003 //      (See accompanying file LICENSE_1_0.txt or copy at
0004 //           https://www.boost.org/LICENSE_1_0.txt)
0005
0006 #ifndef REVERSE_MODE_AUTODIFF_MEMORY_MANAGEMENT_HPP
0007 #define REVERSE_MODE_AUTODIFF_MEMORY_MANAGEMENT_HPP
0008
0009 #include <algorithm>
0010 #include <array>
0011 #include <boost/math/tools/assert.hpp>
0012 #include <cassert>
0013 #include <cstddef>
0014 #include <iterator>
0015 #include <memory>
0016 #include <type_traits>
0017 #include <vector>
0018 namespace boost {
0019 namespace math {
0020 namespace differentiation {
0021 namespace reverse_mode {
0022 namespace detail {
0023 template<typename allocator_type, size_t buffer_size>
0024 class flat_linear_allocator_iterator
0025 {
0026     /**
0027    * @brief enables iterating over linear allocator with
0028    * c++ iterators
0029    */
0030 public:
0031     using raw_allocator_type   = std::remove_const_t<allocator_type>;
0032     using value_type           = typename allocator_type::value_type;
0033     using pointer              = typename allocator_type::value_type *;
0034     using const_ptr_type       = const value_type *;
0035     using reference            = typename allocator_type::value_type &;
0036     using const_reference_type = const value_type &;
0037     using iterator_category    = std::random_access_iterator_tag;
0038     using difference_type      = ptrdiff_t;
0039
0040 private:
0041     const allocator_type *storage_ = nullptr;
0042     size_t                index_   = 0;
0043     size_t                begin_   = 0;
0044     size_t                end_     = 0;
0045
0046 public:
0047     flat_linear_allocator_iterator() = default;
0048
0049     explicit flat_linear_allocator_iterator(allocator_type *storage, size_t index)
0050         : storage_(storage)
0051         , index_(index)
0052         , begin_(0)
0053         , end_(storage->size())
0054     {}
0055
0056     explicit flat_linear_allocator_iterator(allocator_type *storage,
0057                                             size_t          index,
0058                                             size_t          begin,
0059                                             size_t          end)
0060         : storage_(storage)
0061         , index_(index)
0062         , begin_(begin)
0063         , end_(end)
0064     {}
0065
0066     explicit flat_linear_allocator_iterator(const allocator_type *storage, size_t index)
0067         : storage_(storage)
0068         , index_(index)
0069         , begin_(0)
0070         , end_(storage->size())
0071     {}
0072
0073     explicit flat_linear_allocator_iterator(const allocator_type *storage,
0074                                             size_t                index,
0075                                             size_t                begin,
0076                                             size_t                end)
0077         : storage_(storage)
0078         , index_(index)
0079         , begin_(begin)
0080         , end_(end)
0081     {}
0082     reference operator*()
0083     {
0084         BOOST_MATH_ASSERT(index_ >= begin_ && index_ < end_);
0085         return (*storage_->data_[index_ / buffer_size])[index_ % buffer_size];
0086     }
0087
0088     const_reference_type operator*() const
0089     {
0090         BOOST_MATH_ASSERT(index_ >= begin_ && index_ < end_);
0091         return (*storage_->data_[index_ / buffer_size])[index_ % buffer_size];
0092     }
0093
0094     pointer operator->()
0095     {
0096         BOOST_MATH_ASSERT(index_ >= begin_ && index_ < end_);
0097         return &operator*();
0098     }
0099
0100     const_ptr_type operator->() const
0101     {
0102         BOOST_MATH_ASSERT(index_ >= begin_ && index_ < end_);
0103         return &operator*();
0104     }
0105     flat_linear_allocator_iterator &operator++()
0106     {
0107         ++index_;
0108         return *this;
0109     }
0110
0111     flat_linear_allocator_iterator operator++(int)
0112     {
0113         auto tmp = *this;
0114         ++(*this);
0115         return tmp;
0116     }
0117
0118     flat_linear_allocator_iterator &operator--()
0119     {
0120         --index_;
0121         return *this;
0122     }
0123
0124     flat_linear_allocator_iterator operator--(int)
0125     {
0126         auto tmp = *this;
0127         --(*this);
0128         return tmp;
0129     }
0130
0131     bool operator==(const flat_linear_allocator_iterator &other) const
0132     {
0133         return index_ == other.index_ && storage_ == other.storage_;
0134     }
0135
0136     bool operator!=(const flat_linear_allocator_iterator &other) const { return !(*this == other); }
0137
0138     flat_linear_allocator_iterator operator+(difference_type n) const
0139     {
0140         return flat_linear_allocator_iterator(storage_, index_ + static_cast<size_t>(n), begin_, end_);
0141     }
0142
0143     flat_linear_allocator_iterator &operator+=(difference_type n)
0144     {
0145         index_ += n;
0146         return *this;
0147     }
0148
0149     flat_linear_allocator_iterator operator-(difference_type n) const
0150     {
0151         return flat_linear_allocator_iterator(storage_, index_ - n, begin_, end_);
0152     }
0153     flat_linear_allocator_iterator &operator-=(difference_type n)
0154     {
0155         index_ -= n;
0156         return *this;
0157     }
0158
0159     difference_type operator-(const flat_linear_allocator_iterator &other) const
0160     {
0161         return static_cast<difference_type>(index_) - static_cast<difference_type>(other.index_);
0162     }
0163
0164     reference operator[](difference_type n) { return *(*this + n); }
0165
0166     const_reference_type operator[](difference_type n) const { return *(*this + n); }
0167
0168     bool operator<(const flat_linear_allocator_iterator &other) const
0169     {
0170         return index_ < other.index_;
0171     }
0172
0173     bool operator>(const flat_linear_allocator_iterator &other) const
0174     {
0175         return index_ > other.index_;
0176     }
0177
0178     bool operator<=(const flat_linear_allocator_iterator &other) const
0179     {
0180         return index_ <= other.index_;
0181     }
0182
0183     bool operator>=(const flat_linear_allocator_iterator &other) const
0184     {
0185         return index_ >= other.index_;
0186     }
0187
0188     bool operator!() const noexcept { return storage_ == nullptr; }
0189 };
0190 /* memory management helps for tape */
0191 template<typename RealType, size_t buffer_size>
0192 class flat_linear_allocator
0193 {
0194     /** @brief basically a vector<array<T*, size>>
0195    * intended to work like a vector that allocates memory in chunks
0196    * and doesn't invalidate references
0197    * */
0198 public:
0199     // store vector of unique pointers to arrays
0200     // to avoid vector reference invalidation
0201     using buffer_type = std::array<RealType, buffer_size>;
0202     using buffer_ptr  = std::unique_ptr<std::array<RealType, buffer_size>>;
0203
0204 private:
0205     std::vector<buffer_ptr> data_;
0206     size_t                  total_size_ = 0;
0207     std::vector<size_t>     checkpoints_; //{0};
0208
0209 public:
0210     friend class flat_linear_allocator_iterator<flat_linear_allocator<RealType, buffer_size>,
0211                                                 buffer_size>;
0212     friend class flat_linear_allocator_iterator<const flat_linear_allocator<RealType, buffer_size>,
0213                                                 buffer_size>;
0214     using value_type = RealType;
0215     using iterator
0216         = flat_linear_allocator_iterator<flat_linear_allocator<RealType, buffer_size>, buffer_size>;
0217     using const_iterator
0218         = flat_linear_allocator_iterator<const flat_linear_allocator<RealType, buffer_size>,
0219                                          buffer_size>;
0220
0221     size_t buffer_id() const noexcept { return total_size_ / buffer_size; }
0222     size_t item_id() const noexcept { return total_size_ % buffer_size; }
0223
0224 private:
0225     void allocate_buffer()
0226     {
0227         data_.emplace_back(std::make_unique<buffer_type>());
0228     }
0229
0230 public:
0231     flat_linear_allocator() { allocate_buffer(); }
0232     flat_linear_allocator(const flat_linear_allocator &)            = delete;
0233     flat_linear_allocator &operator=(const flat_linear_allocator &) = delete;
0234     flat_linear_allocator(flat_linear_allocator &&)                 = delete;
0235     flat_linear_allocator &operator=(flat_linear_allocator &&)      = delete;
0236     ~flat_linear_allocator()
0237     {
0238         destroy_all();
0239         data_.clear();
0240     }
0241
0242     void destroy_all()
0243     {
0244         for (size_t i = 0; i < total_size_; ++i) {
0245             size_t bid = i / buffer_size;
0246             size_t iid = i % buffer_size;
0247             (*data_[bid])[iid].~RealType();
0248         }
0249     }
0250     /** @brief
0251    * helper functions to clear tape and create block in tape
0252    */
0253     void clear()
0254     {
0255         data_.clear();
0256         total_size_ = 0;
0257         checkpoints_.clear();
0258         allocate_buffer();
0259     }
0260
0261     // doesn't delete anything, only sets the current index to zero
0262     void reset() { total_size_ = 0; }
0263     void rewind() { total_size_ = 0; };
0264
0265     // adds current index as a checkpoint to be able to walk back to
0266     void add_checkpoint()
0267     {
0268         if (total_size_ > 0) {
0269             checkpoints_.push_back(total_size_ - 1);
0270         } else {
0271             checkpoints_.push_back(0);
0272         }
0273     };
0274
0275     /** @brief clears all checkpoints
0276    * */
0277     void reset_checkpoints() { checkpoints_.clear(); }
0278
0279     void rewind_to_last_checkpoint() { total_size_ = checkpoints_.back(); }
0280     void rewind_to_checkpoint_at(size_t index) { total_size_ = checkpoints_[index]; }
0281
0282     void fill(const RealType &val)
0283     {
0284         for (size_t i = 0; i < total_size_; ++i) {
0285             size_t bid         = i / buffer_size;
0286             size_t iid         = i % buffer_size;
0287             (*data_[bid])[iid] = val;
0288         }
0289     }
0290
0291     /** @brief emplaces back object at the end of the
0292    * data structure, calls default constructor */
0293     iterator emplace_back()
0294     {
0295         if (item_id() == 0 && total_size_ != 0) {
0296             allocate_buffer();
0297         }
0298         size_t bid = buffer_id();
0299         size_t iid = item_id();
0300
0301         RealType *ptr = &(*data_[bid])[iid];
0302         new (ptr) RealType();
0303         ++total_size_;
0304         return iterator(this, total_size_ - 1);
0305     };
0306
0307     /** @brief, emplaces back object at end of data structure,
0308    * passes arguments to constructor */
0309     template<typename... Args>
0310     iterator emplace_back(Args &&...args)
0311     {
0312         if (item_id() == 0 && total_size_ != 0) {
0313             allocate_buffer();
0314         }
0315         BOOST_MATH_ASSERT(buffer_id() < data_.size());
0316         BOOST_MATH_ASSERT(item_id() < buffer_size);
0317         RealType *ptr = &(*data_[buffer_id()])[item_id()];
0318         new (ptr) RealType(std::forward<Args>(args)...);
0319         ++total_size_;
0320         return iterator(this, total_size_ - 1);
0321     }
0322     /** @brief default constructs n objects at end of
0323    * data structure, n known at compile time */
0324     template<size_t n>
0325     iterator emplace_back_n()
0326     {
0327         size_t bid = buffer_id();
0328         size_t iid = item_id();
0329         if (iid + n < buffer_size) {
0330             RealType *ptr = &(*data_[bid])[iid];
0331             for (size_t i = 0; i < n; ++i) {
0332                 new (ptr + i) RealType();
0333             }
0334             total_size_ += n;
0335             return iterator(this, total_size_ - n, total_size_ - n, total_size_);
0336         } else {
0337             size_t allocs_in_curr_buffer = buffer_size - iid;
0338             size_t allocs_in_next_buffer = n - (buffer_size - iid);
0339             RealType *ptr                   = &(*data_[bid])[iid];
0340             for (size_t i = 0; i < allocs_in_curr_buffer; ++i) {
0341                 new (ptr + i) RealType();
0342             }
0343             allocate_buffer();
0344             bid = data_.size() - 1;
0345             iid = 0;
0346             total_size_ += n;
0347
0348             RealType *ptr2 = &(*data_[bid])[iid];
0349             for (size_t i = 0; i < allocs_in_next_buffer; i++) {
0350                 new (ptr2 + i) RealType();
0351             }
0352             return iterator(this, total_size_ - n, total_size_ - n, total_size_);
0353         }
0354     }
0355     /** @brief default constructs n objects at end of
0356    * data structure, n known at run time
0357    */
0358     iterator emplace_back_n(size_t n)
0359     {
0360         size_t bid = buffer_id();
0361         size_t iid = item_id();
0362         if (iid + n < buffer_size) {
0363             RealType *ptr = &(*data_[bid])[iid];
0364             for (size_t i = 0; i < n; ++i) {
0365                 new (ptr + i) RealType();
0366             }
0367             total_size_ += n;
0368             return iterator(this, total_size_ - n, total_size_ - n, total_size_);
0369         } else {
0370             size_t allocs_in_curr_buffer = buffer_size - iid;
0371             size_t allocs_in_next_buffer = n - (buffer_size - iid);
0372             RealType *ptr                   = &(*data_[bid])[iid];
0373             for (size_t i = 0; i < allocs_in_curr_buffer; ++i) {
0374                 new (ptr + i) RealType();
0375             }
0376             allocate_buffer();
0377             bid = data_.size() - 1;
0378             iid = 0;
0379             total_size_ += n;
0380
0381             RealType *ptr2 = &(*data_[bid])[iid];
0382             for (size_t i = 0; i < allocs_in_next_buffer; i++) {
0383                 new (ptr2 + i) RealType();
0384             }
0385             return iterator(this, total_size_ - n, total_size_ - n, total_size_);
0386         }
0387     }
0388
0389     /** @brief number of elements */
0390     size_t size() const { return total_size_; }
0391
0392     /** @brief total capacity */
0393     size_t capacity() const { return data_.size() * buffer_size; }
0394
0395     /** @brief iterator helpers */
0396     iterator       begin() { return iterator(this, 0); }
0397     iterator       end() { return iterator(this, total_size_); }
0398     const_iterator begin() const { return const_iterator(this, 0); }
0399     const_iterator end() const { return const_iterator(this, total_size_); }
0400
0401     iterator last_checkpoint() { return iterator(this, checkpoints_.back(), 0, total_size_); }
0402     iterator first_checkpoint() { return iterator(this, checkpoints_[0], 0, total_size_); };
0403     iterator checkpoint_at(size_t index)
0404     {
0405         return iterator(this, checkpoints_[index], 0, total_size_);
0406     };
0407
0408     /** @brief searches for item in allocator
0409    *  only used to find gradient nodes for propagation */
0410     iterator find(const RealType *const item)
0411     {
0412         return std::find_if(begin(), end(), [&](const RealType &val) { return &val == item; });
0413     }
0414     /** @brief vector like access,
0415    *  currently unused anywhere but very useful for debugging
0416    */
0417     RealType &operator[](std::size_t i)
0418     {
0419         BOOST_MATH_ASSERT(i < total_size_);
0420         return (*data_[i / buffer_size])[i % buffer_size];
0421     }
0422     const RealType &operator[](std::size_t i) const
0423     {
0424         BOOST_MATH_ASSERT(i < total_size_);
0425         return (*data_[i / buffer_size])[i % buffer_size];
0426     }
0427 };
0428 } // namespace detail
0429 } // namespace reverse_mode
0430 } // namespace differentiation
0431 } // namespace math
0432 } // namespace boost
0433
0434 #endif