From edf7593e33a822894663d4da12efb05a1e21e9e3 Mon Sep 17 00:00:00 2001
From: Vaibhav Thakkar <vaibhav.thakkar.22.12.99@gmail.com>
Date: Thu, 20 Jun 2024 21:58:38 +0200
Subject: [PATCH] Add benchmark for hessians

---
 benchmark/CMakeLists.txt |  1 +
 benchmark/Hessians.cpp   | 47 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)
 create mode 100644 benchmark/Hessians.cpp

diff --git a/benchmark/CMakeLists.txt b/benchmark/CMakeLists.txt
index d0559faf2..9ae784607 100644
--- a/benchmark/CMakeLists.txt
+++ b/benchmark/CMakeLists.txt
@@ -9,6 +9,7 @@ if (CLAD_ENABLE_ENZYME_BACKEND)
 endif(CLAD_ENABLE_ENZYME_BACKEND)
 CB_ADD_GBENCHMARK(VectorModeComparison VectorModeComparison.cpp)
 CB_ADD_GBENCHMARK(MemoryComplexity MemoryComplexity.cpp)
+CB_ADD_GBENCHMARK(Hessians Hessians.cpp)
 
 set (CLAD_BENCHMARK_DEPS clad)
 get_property(_benchmark_names DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} PROPERTY TESTS)
diff --git a/benchmark/Hessians.cpp b/benchmark/Hessians.cpp
new file mode 100644
index 000000000..ae3a28a33
--- /dev/null
+++ b/benchmark/Hessians.cpp
@@ -0,0 +1,47 @@
+#include "benchmark/benchmark.h"
+
+#include "clad/Differentiator/Differentiator.h"
+
+#include "BenchmarkedFunctions.h"
+
+// Benchmark Hessian diagonal sum computation, by computing the
+// entire computation.
+static void BM_HessianCompleteComputation(benchmark::State& state) {
+  auto dfdx2 = clad::hessian(weightedSum, "p[0:1],w[0:1]");
+  double p[] = {1, 2};
+  double w[] = {3, 4};
+  unsigned long long sum = 0;
+  double hessianMatrix[16] = {};
+  for (unsigned i = 0; i < 16; i++)
+    hessianMatrix[i] = 0.0;
+  for (auto _ : state) {
+    dfdx2.execute(p, w, 3, hessianMatrix);
+    for (int i = 0; i < 4; i++)
+      // Sum the diagonal of the Hessian matrix.
+      benchmark::DoNotOptimize(sum += hessianMatrix[i * 4 + i]);
+  }
+}
+BENCHMARK(BM_HessianCompleteComputation);
+
+// Benchmark Hessian diagonal sum computation, by computing only
+// the diagonal elements.
+static void BM_HessianDiagonalComputation(benchmark::State& state) {
+  auto dfdx2 =
+      clad::hessian<clad::opts::diagonal_only>(weightedSum, "p[0:1],w[0:1]");
+  double p[] = {1, 2};
+  double w[] = {3, 4};
+  unsigned long long sum = 0;
+  double diagonalHessian[4] = {};
+  for (unsigned i = 0; i < 4; i++)
+    diagonalHessian[i] = 0.0;
+  for (auto _ : state) {
+    dfdx2.execute(p, w, 3, diagonalHessian);
+    for (int i = 0; i < 4; i++)
+      // Sum the diagonal of the Hessian matrix.
+      benchmark::DoNotOptimize(sum += diagonalHessian[i]);
+  }
+}
+BENCHMARK(BM_HessianDiagonalComputation);
+
+// Define our main.
+BENCHMARK_MAIN();