From 9537d65fa57602146b193bc85bde3c1eb5d01ffb Mon Sep 17 00:00:00 2001
From: Tyler Veness <calcmogul@gmail.com>
Date: Fri, 17 Jan 2025 13:03:06 -0800
Subject: [PATCH] Include gradient tree generation in Hessian setup time (#691)

---
 include/sleipnir/autodiff/Hessian.hpp  | 16 +++++++++++++---
 include/sleipnir/autodiff/Jacobian.hpp |  5 +++--
 jormungandr/cpp/Docstrings.hpp         |  2 ++
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/include/sleipnir/autodiff/Hessian.hpp b/include/sleipnir/autodiff/Hessian.hpp
index 4563aa23..60dd3b66 100644
--- a/include/sleipnir/autodiff/Hessian.hpp
+++ b/include/sleipnir/autodiff/Hessian.hpp
@@ -36,6 +36,8 @@ class SLEIPNIR_DLLEXPORT Hessian {
   Hessian(Variable variable, const VariableMatrix& wrt) noexcept
       : m_jacobian{
             [&] {
+              m_profiler.StartSetup();
+
               small_vector<detail::ExpressionPtr> wrtVec;
               wrtVec.reserve(wrt.size());
               for (auto& elem : wrt) {
@@ -52,7 +54,9 @@ class SLEIPNIR_DLLEXPORT Hessian {
               }
               return ret;
             }(),
-            wrt} {}
+            wrt} {
+    m_profiler.StopSetup();
+  }
 
   /**
    * Returns the Hessian as a VariableMatrix.
@@ -65,14 +69,20 @@ class SLEIPNIR_DLLEXPORT Hessian {
   /**
    * Evaluates the Hessian at wrt's value.
    */
-  const Eigen::SparseMatrix<double>& Value() { return m_jacobian.Value(); }
+  const Eigen::SparseMatrix<double>& Value() {
+    m_profiler.StartSolve();
+    const auto& H = m_jacobian.Value();
+    m_profiler.StopSolve();
+    return H;
+  }
 
   /**
    * Returns the profiler.
    */
-  Profiler& GetProfiler() { return m_jacobian.GetProfiler(); }
+  Profiler& GetProfiler() { return m_profiler; }
 
  private:
+  Profiler m_profiler;
   Jacobian m_jacobian;
 };
 
diff --git a/include/sleipnir/autodiff/Jacobian.hpp b/include/sleipnir/autodiff/Jacobian.hpp
index 47a64654..6006db96 100644
--- a/include/sleipnir/autodiff/Jacobian.hpp
+++ b/include/sleipnir/autodiff/Jacobian.hpp
@@ -98,12 +98,13 @@ class SLEIPNIR_DLLEXPORT Jacobian {
    * Evaluates the Jacobian at wrt's value.
    */
   const Eigen::SparseMatrix<double>& Value() {
+    m_profiler.StartSolve();
+
     if (m_nonlinearRows.empty()) {
+      m_profiler.StopSolve();
       return m_J;
     }
 
-    m_profiler.StartSolve();
-
     for (auto& graph : m_graphs) {
       graph.Update();
     }
diff --git a/jormungandr/cpp/Docstrings.hpp b/jormungandr/cpp/Docstrings.hpp
index 242abf0d..e118e79f 100644
--- a/jormungandr/cpp/Docstrings.hpp
+++ b/jormungandr/cpp/Docstrings.hpp
@@ -203,6 +203,8 @@ static const char *__doc_sleipnir_Hessian_Value = R"doc(Evaluates the Hessian at
 
 static const char *__doc_sleipnir_Hessian_m_jacobian = R"doc()doc";
 
+static const char *__doc_sleipnir_Hessian_m_profiler = R"doc()doc";
+
 static const char *__doc_sleipnir_InequalityConstraints = R"doc(A vector of inequality constraints of the form cᵢ(x) ≥ 0.)doc";
 
 static const char *__doc_sleipnir_InequalityConstraints_InequalityConstraints =