Stop pre-constructing reverse-mode tape blocks, and reuse them after a rewind

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<rvar<double, 1>> x(n, rvar<double, 1>(0))` holds one variable n times. Write `std::vector<rvar<double, 1>> 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]]

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<RealType, buffer_size>;
-  using buffer_ptr = std::unique_ptr<std::array<RealType, buffer_size>>;
+  // Uninitialized storage: elements are constructed by emplace_back. Default-constructing a whole block
+  // would, for the rvar<T, 1> 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<RealType*>(bytes + sizeof(RealType) * i); }
+    const RealType& operator[](size_t i) const noexcept { return *reinterpret_cast<const RealType*>(bytes + sizeof(RealType) * i); }
+  };
+  using buffer_ptr = std::unique_ptr<buffer_type>;
 
 private:
   std::vector<buffer_ptr> 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<buffer_type>());
+    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<typename... Args>
   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>(args)...);
+    new (slot(total_size_)) RealType(std::forward<Args>(args)...);
     ++total_size_;
     return iterator(this, total_size_ - 1);
   }
@@ -348,66 +346,18 @@ class flat_linear_allocator
   template<size_t n>
   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 */

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 <class T>
+void flat_linear_allocator_test_block_reuse()
+{
+  constexpr size_t buffer_size = 16;
+  rdiff_detail::flat_linear_allocator<T, buffer_size> 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<counted, 16> 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 <class T>
+void second_order_tape_adds_no_first_order_nodes()
+{
+  auto& tape1 = rdiff::get_active_tape<T, 1>();
+  auto before = tape1.end() - tape1.begin();
+  (void)rdiff::get_active_tape<T, 2>();
+  CHECK_EQUAL(tape1.end() - tape1.begin(), before);
+}
+
 int main()
 {
    boost::mp11::mp_for_each<boost::mp11::mp_transform<boost::mp11::mp_identity, all_float_types>>([](auto t) { flat_linear_allocator_constructors<typename decltype(t)::type>(); });
    boost::mp11::mp_for_each<boost::mp11::mp_transform<boost::mp11::mp_identity, all_float_types>>([](auto t) { flat_linear_allocator_test_emplace<typename decltype(t)::type>(); });
    boost::mp11::mp_for_each<boost::mp11::mp_transform<boost::mp11::mp_identity, all_float_types>>([](auto t) { flat_linear_allocator_test_checkpointing<typename decltype(t)::type>(); });
+   boost::mp11::mp_for_each<boost::mp11::mp_transform<boost::mp11::mp_identity, boost::mp11::mp_list<float, double, cpp_bin_float_50>>>([](auto t) { flat_linear_allocator_test_block_reuse<typename decltype(t)::type>(); });
+   flat_linear_allocator_test_destruction();
+   boost::mp11::mp_for_each<boost::mp11::mp_transform<boost::mp11::mp_identity, boost::mp11::mp_list<float, double, cpp_bin_float_50>>>([](auto t) { second_order_tape_adds_no_first_order_nodes<typename decltype(t)::type>(); });
    return boost::math::test::report_errors();
 }