Fix adjoint memory leak in ExpressionGraph

There was a circular reference between adjoints and their parent
expressions that prevented them from being freed. The adjoints are now
unlinked when they're no longer needed.

The Hessian tests were instrumented to confirm all Expressions were
freed at the end of each test.

PoolResource and PoolAllocator were restructured to take the chunk size
at runtime for ease of use, and the default chunk size was reduced to
16384.

include/sleipnir/autodiff/Expression.hpp

@@ -1154,20 +1154,3 @@ SLEIPNIR_DLLEXPORT inline ExpressionPtr tanh(const ExpressionPtr& x) {
 }
 
 }  // namespace sleipnir::detail
-
-namespace sleipnir {
-
-// FIXME: Doxygen is confused:
-//
-//   Found ';' while parsing initializer list! (doxygen could be confused by a
-//   macro call without semicolon)
-
-//! @cond Doxygen_Suppress
-
-// Instantiate Expression pool in Expression.cpp instead to avoid ODR violation
-extern template EXPORT_TEMPLATE_DECLARE(SLEIPNIR_DLLEXPORT)
-    PoolAllocator<detail::Expression> GlobalPoolAllocator<detail::Expression>();
-
-//! @endcond
-
-}  // namespace sleipnir

include/sleipnir/autodiff/ExpressionGraph.hpp

@@ -167,6 +167,17 @@ class SLEIPNIR_DLLEXPORT ExpressionGraph {
       }
     }
 
+    // Unlink adjoints to avoid circular references between them and their
+    // parent expressions. This ensures all expressions are returned to the free
+    // list.
+    for (auto node : m_adjointList) {
+      for (auto& arg : node->args) {
+        if (arg != nullptr) {
+          arg->adjointExpr = nullptr;
+        }
+      }
+    }
+
     for (size_t row = 0; row < wrt.size(); ++row) {
       wrt[row]->row = -1;
     }

include/sleipnir/util/Pool.hpp

@@ -14,18 +14,21 @@ namespace sleipnir {
  * The pool allocates chunks of memory and splits them into blocks managed by a
  * free list. Allocations return pointers from the free list, and deallocations
  * return pointers to the free list.
- *
- * @tparam BlocksPerChunk Number of blocks per chunk of memory.
  */
-template <size_t BlocksPerChunk>
 class PoolResource {
  public:
-  constexpr PoolResource() = default;
+  /**
+   * Constructs a default PoolResource.
+   *
+   * @param blocksPerChunk Number of blocks per chunk of memory.
+   */
+  explicit PoolResource(size_t blocksPerChunk)
+      : blocksPerChunk{blocksPerChunk} {}
 
-  constexpr PoolResource(const PoolResource&) = delete;
-  constexpr PoolResource& operator=(const PoolResource&) = delete;
-  constexpr PoolResource(PoolResource&&) = default;
-  constexpr PoolResource& operator=(PoolResource&&) = default;
+  PoolResource(const PoolResource&) = delete;
+  PoolResource& operator=(const PoolResource&) = delete;
+  PoolResource(PoolResource&&) = default;
+  PoolResource& operator=(PoolResource&&) = default;
 
   /**
    * Returns a block of memory from the pool.
@@ -60,14 +63,21 @@ class PoolResource {
   /**
    * Returns true if this pool resource has the same backing storage as another.
    */
-  constexpr bool is_equal(
-      const PoolResource<BlocksPerChunk>& other) const noexcept {
+  constexpr bool is_equal(const PoolResource& other) const noexcept {
     return this == &other;
   }
 
+  /**
+   * Returns the number of blocks from this pool resource that are in use.
+   */
+  size_t blocks_in_use() const noexcept {
+    return m_buffer.size() * blocksPerChunk - m_freeList.size();
+  }
+
  private:
   std::vector<std::unique_ptr<std::byte[]>> m_buffer;
   std::vector<void*> m_freeList;
+  size_t blocksPerChunk;
 
   /**
    * Adds a memory chunk to the pool, partitions it into blocks with the given
@@ -76,8 +86,8 @@ class PoolResource {
    * @param bytesPerBlock Number of bytes in the block.
    */
   constexpr void AddChunk(size_t bytesPerBlock) {
-    m_buffer.emplace_back(new std::byte[bytesPerBlock * BlocksPerChunk]);
-    for (int i = BlocksPerChunk - 1; i >= 0; --i) {
+    m_buffer.emplace_back(new std::byte[bytesPerBlock * blocksPerChunk]);
+    for (int i = blocksPerChunk - 1; i >= 0; --i) {
       m_freeList.emplace_back(m_buffer.back().get() + bytesPerBlock * i);
     }
   }
@@ -87,9 +97,8 @@ class PoolResource {
  * This class is an allocator for the pool resource.
  *
  * @tparam T The type of object in the pool.
- * @tparam ObjectsPerChunk Number of objects per chunk of memory.
  */
-template <typename T, size_t ObjectsPerChunk = 32768>
+template <typename T>
 class PoolAllocator {
  public:
   /**
@@ -102,17 +111,14 @@ class PoolAllocator {
    *
    * @param r The pool resource.
    */
-  explicit constexpr PoolAllocator(PoolResource<ObjectsPerChunk>* r)
-      : m_memoryResource{r} {}
+  explicit constexpr PoolAllocator(PoolResource* r) : m_memoryResource{r} {}
 
   /**
    * Copy constructor.
    */
-  constexpr PoolAllocator(const PoolAllocator<T, ObjectsPerChunk>& other) =
-      default;
+  constexpr PoolAllocator(const PoolAllocator<T>& other) = default;
 
-  constexpr PoolAllocator<T>& operator=(
-      const PoolAllocator<T, ObjectsPerChunk>&) = delete;
+  constexpr PoolAllocator<T>& operator=(const PoolAllocator<T>&) = delete;
 
   /**
    * Returns a block of memory from the pool.
@@ -135,19 +141,22 @@ class PoolAllocator {
   }
 
  private:
-  PoolResource<ObjectsPerChunk>* m_memoryResource;
+  PoolResource* m_memoryResource;
 };
 
+/**
+ * Returns a global pool memory resource.
+ */
+PoolResource& GlobalPoolResource();
+
 /**
  * Returns an allocator for a global pool memory resource.
  *
  * @tparam T The type of object in the pool.
- * @tparam ObjectsPerChunk Number of objects per chunk of memory.
  */
-template <typename T, size_t ObjectsPerChunk = 32768>
-PoolAllocator<T, ObjectsPerChunk> GlobalPoolAllocator() {
-  static PoolResource<ObjectsPerChunk> pool;
-  return PoolAllocator<T, ObjectsPerChunk>{&pool};
+template <typename T>
+PoolAllocator<T> GlobalPoolAllocator() {
+  return PoolAllocator<T>{&GlobalPoolResource()};
 }
 
 }  // namespace sleipnir

src/util/Pool.cpp

@@ -0,0 +1,12 @@
+// Copyright (c) Sleipnir contributors
+
+#include "sleipnir/util/Pool.hpp"
+
+namespace sleipnir {
+
+PoolResource& GlobalPoolResource() {
+  static PoolResource pool{16384};
+  return pool;
+}
+
+}  // namespace sleipnir

test/src/autodiff/HessianTest.cpp

@@ -12,8 +12,12 @@
 #include <sleipnir/autodiff/Variable.hpp>
 
 #include "Range.hpp"
+#include "util/ScopeExit.hpp"
 
 TEST_CASE("Hessian - Linear", "[Hessian]") {
+  sleipnir::scope_exit exit{
+      [] { CHECK(sleipnir::GlobalPoolResource().blocks_in_use() == 0u); }};
+
   // y = x
   sleipnir::VariableMatrix x{1};
   x(0).SetValue(3);
@@ -30,6 +34,9 @@ TEST_CASE("Hessian - Linear", "[Hessian]") {
 }
 
 TEST_CASE("Hessian - Quadratic", "[Hessian]") {
+  sleipnir::scope_exit exit{
+      [] { CHECK(sleipnir::GlobalPoolResource().blocks_in_use() == 0u); }};
+
   // y = x²
   // y = x * x
   sleipnir::VariableMatrix x{1};
@@ -51,6 +58,9 @@ TEST_CASE("Hessian - Quadratic", "[Hessian]") {
 }
 
 TEST_CASE("Hessian - Sum", "[Hessian]") {
+  sleipnir::scope_exit exit{
+      [] { CHECK(sleipnir::GlobalPoolResource().blocks_in_use() == 0u); }};
+
   sleipnir::VariableMatrix x{5};
   x(0).SetValue(1);
   x(1).SetValue(2);
@@ -72,6 +82,9 @@ TEST_CASE("Hessian - Sum", "[Hessian]") {
 }
 
 TEST_CASE("Hessian - Sum of products", "[Hessian]") {
+  sleipnir::scope_exit exit{
+      [] { CHECK(sleipnir::GlobalPoolResource().blocks_in_use() == 0u); }};
+
   sleipnir::VariableMatrix x{5};
   x(0).SetValue(1);
   x(1).SetValue(2);
@@ -95,6 +108,9 @@ TEST_CASE("Hessian - Sum of products", "[Hessian]") {
 }
 
 TEST_CASE("Hessian - Product of sines", "[Hessian]") {
+  sleipnir::scope_exit exit{
+      [] { CHECK(sleipnir::GlobalPoolResource().blocks_in_use() == 0u); }};
+
   sleipnir::VariableMatrix x{5};
   x(0).SetValue(1);
   x(1).SetValue(2);
@@ -148,6 +164,9 @@ TEST_CASE("Hessian - Product of sines", "[Hessian]") {
 }
 
 TEST_CASE("Hessian - Sum of squared residuals", "[Hessian]") {
+  sleipnir::scope_exit exit{
+      [] { CHECK(sleipnir::GlobalPoolResource().blocks_in_use() == 0u); }};
+
   sleipnir::Variable y;
   Eigen::VectorXd g;
   Eigen::MatrixXd H;
@@ -188,6 +207,9 @@ TEST_CASE("Hessian - Sum of squared residuals", "[Hessian]") {
 }
 
 TEST_CASE("Hessian - Sum of squares", "[Hessian]") {
+  sleipnir::scope_exit exit{
+      [] { CHECK(sleipnir::GlobalPoolResource().blocks_in_use() == 0u); }};
+
   sleipnir::VariableMatrix r{4};
   r(0).SetValue(25.0);
   r(1).SetValue(10.0);
@@ -212,6 +234,9 @@ TEST_CASE("Hessian - Sum of squares", "[Hessian]") {
 }
 
 TEST_CASE("Hessian - Rosenbrock", "[Hessian]") {
+  sleipnir::scope_exit exit{
+      [] { CHECK(sleipnir::GlobalPoolResource().blocks_in_use() == 0u); }};
+
   sleipnir::VariableMatrix input{2};
   auto& x = input(0);
   auto& y = input(1);
@@ -234,6 +259,9 @@ TEST_CASE("Hessian - Rosenbrock", "[Hessian]") {
 }
 
 TEST_CASE("Hessian - Variable reuse", "[Hessian]") {
+  sleipnir::scope_exit exit{
+      [] { CHECK(sleipnir::GlobalPoolResource().blocks_in_use() == 0u); }};
+
   sleipnir::Variable y;
   sleipnir::VariableMatrix x{1};