From 8fd8064a99ce9f654148ab335a939442aa606ded Mon Sep 17 00:00:00 2001
From: Rohan Julka <rohanjulka19@gmail.com>
Date: Thu, 20 Jun 2024 16:32:51 +0100
Subject: [PATCH] Use connected lists for tape data structure

Instead of a single list use mutliple connected lists to store elements.
This allows to dynamically increase the size of tape without the need of
relocating elements.
---
 benchmark/AlgorithmicComplexity.cpp          |  2 +-
 demos/CustomTypeNumDiff.cpp                  |  2 +-
 include/clad/Differentiator/Differentiator.h | 18 +++--
 include/clad/Differentiator/NewTape.h        | 84 ++++++++++++++++++++
 test/NumericalDiff/PureCentralDiffCalls.C    |  2 +-
 5 files changed, 97 insertions(+), 11 deletions(-)
 create mode 100644 include/clad/Differentiator/NewTape.h

diff --git a/benchmark/AlgorithmicComplexity.cpp b/benchmark/AlgorithmicComplexity.cpp
index 9717cb0cd..15089a24e 100644
--- a/benchmark/AlgorithmicComplexity.cpp
+++ b/benchmark/AlgorithmicComplexity.cpp
@@ -19,7 +19,7 @@ static void BM_NumericGausP(benchmark::State& state) {
   double p[] = {1, 2, 3, 4, 5};
   double dx[5] = {0, 0, 0, 0, 0};
   double dp[5] = {0, 0, 0, 0, 0};
-  clad::tape<clad::array_ref<double>> results = {};
+  clad::old_tape<clad::array_ref<double>> results = {};
   int dim = 5;
   results.emplace_back(dx, dim);
   results.emplace_back(dp, dim);
diff --git a/demos/CustomTypeNumDiff.cpp b/demos/CustomTypeNumDiff.cpp
index 0f92c22f9..33bc6adb1 100644
--- a/demos/CustomTypeNumDiff.cpp
+++ b/demos/CustomTypeNumDiff.cpp
@@ -134,7 +134,7 @@ int main() {
   // This is how we return the derivative with respect to all arguments.
   // The order of being placed in this tape should be the same as the order of
   // the arguments being passed to the function.
-  clad::tape<clad::array_ref<
+  clad::old_tape<clad::array_ref<
       double /*This should be the return value of the function you want to differentiate.*/>>
       grad = {};
   // Place the l-value reference of the variables in the tape.
diff --git a/include/clad/Differentiator/Differentiator.h b/include/clad/Differentiator/Differentiator.h
index 4a089c095..d35114fe1 100644
--- a/include/clad/Differentiator/Differentiator.h
+++ b/include/clad/Differentiator/Differentiator.h
@@ -14,6 +14,7 @@
 #include "DynamicGraph.h"
 #include "FunctionTraits.h"
 #include "Matrix.h"
+#include "NewTape.h"
 #include "NumericalDiff.h"
 #include "Tape.h"
 
@@ -43,15 +44,16 @@ inline CUDA_HOST_DEVICE unsigned int GetLength(const char* code) {
 }
 
   /// Tape type used for storing values in reverse-mode AD inside loops.
-  template <typename T>
-  using tape = tape_impl<T>;
+template <typename T> using tape = new_tape_impl<T>;
 
-  /// Add value to the end of the tape, return the same value.
-  template <typename T, typename... ArgsT>
-  CUDA_HOST_DEVICE T push(tape<T>& to, ArgsT... val) {
-    to.emplace_back(std::forward<ArgsT>(val)...);
-    return to.back();
-  }
+template <typename T> using old_tape = tape_impl<T>;
+
+/// Add value to the end of the tape, return the same value.
+template <typename T, typename... ArgsT>
+CUDA_HOST_DEVICE T push(tape<T>& to, ArgsT... val) {
+  to.emplace_back(std::forward<ArgsT>(val)...);
+  return to.back();
+}
 
   /// Add value to the end of the tape, return the same value.
   /// A specialization for clad::array_ref types to use in reverse mode.
diff --git a/include/clad/Differentiator/NewTape.h b/include/clad/Differentiator/NewTape.h
new file mode 100644
index 000000000..3d3a3ac18
--- /dev/null
+++ b/include/clad/Differentiator/NewTape.h
@@ -0,0 +1,84 @@
+#ifndef CLAD_NEWTAPE_H
+#define CLAD_NEWTAPE_H
+
+#include <cassert>
+#include <cstdio>
+#include <memory>
+#include <type_traits>
+#include <utility>
+
+namespace clad {
+
+template <typename T> class Block {
+public:
+  T* _data;
+  Block<T>* next;
+  Block<T>* prev;
+
+  Block(std::size_t& _capacity) {
+    _data =
+        static_cast<T*>(::operator new(_capacity * sizeof(T), std::nothrow));
+    next = nullptr;
+    prev = nullptr;
+  }
+
+  ~Block() {}
+};
+
+template <typename T> class new_tape_impl {
+  Block<T>* cur_block = nullptr;
+  std::size_t _capacity = 32;
+  std::size_t _size = 0;
+  using pointer = T*;
+  using reference = T&;
+
+  using iterator = pointer;
+
+public:
+  template <typename... ArgsT>
+
+  void emplace_back(ArgsT&&... args) {
+    if (!cur_block || _size >= _capacity) {
+      Block<T>* prev_block = cur_block;
+      cur_block = new Block<T>(_capacity);
+      if (prev_block) {
+        prev_block->next = cur_block;
+        cur_block->prev = prev_block;
+      }
+      _size = 0;
+    }
+    _size += 1;
+    ::new (const_cast<void*>(static_cast<const volatile void*>(block_end())))
+        T(std::forward<ArgsT>(args)...);
+  }
+
+  std::size_t size() const { return _size; }
+
+  iterator block_begin() {
+    return reinterpret_cast<iterator>(cur_block->_data);
+  }
+
+  iterator block_end() {
+    return reinterpret_cast<iterator>(cur_block->_data + _size - 1);
+  }
+
+  reference back() {
+    assert(_size || cur_block->prev);
+    return block_begin()[_size - 1];
+  }
+
+  void pop_back() {
+    assert(_size || cur_block->prev);
+    block_end()->~T();
+    _size -= 1;
+    if (_size == 0) {
+      Block<T>* temp = cur_block;
+      cur_block = cur_block->prev;
+      temp->~Block<T>();
+      _size = _capacity;
+    }
+  }
+};
+} // namespace clad
+
+#endif // CLAD_NEWTAPE_H
diff --git a/test/NumericalDiff/PureCentralDiffCalls.C b/test/NumericalDiff/PureCentralDiffCalls.C
index b96b13515..fdbbb582b 100644
--- a/test/NumericalDiff/PureCentralDiffCalls.C
+++ b/test/NumericalDiff/PureCentralDiffCalls.C
@@ -75,7 +75,7 @@ int main() { // expected-no-diagnostics
   printf("Result is = %f\n", func1_res); // CHECK-EXEC: Result is = 2.000000
 
   // Gradients, derivative wrt all args
-  clad::tape<clad::array_ref<double>> grad = {};
+  clad::old_tape<clad::array_ref<double>> grad = {};
   grad.emplace_back(dx, 3);
   grad.emplace_back(&dy);
   grad.emplace_back(&dz);