benchmark.cpp

// benchmark.cpp
//
// Standalone performance harness for the HNSW index. Defines its own main(),
// so it is built and run separately from the test suite (test.cpp).
//
// Build & run:
//   make bench && ./benchout
//
// Reports, at a few dataset sizes:
//   - insert throughput (vectors/sec)
//   - query latency (avg / p50 / p95 / p99, ms)
//   - recall@k vs. a brute-force cosine-similarity ground truth
//
// The point is to re-run this after each optimization pass (better graph
// construction, smarter search pruning, caching, etc.) and compare the
// numbers against the previous run.

#include "Embedding.hpp"
#include "Hnsw.hpp"
#include "mem_probe.hpp"
#include <algorithm>
#include <chrono>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <numeric>
#include <random>
#include <string>
#include <vector>
// NEW: replaces global new/delete process-wide.
// Include from this TU only — never test.cpp.
using memprobe::human;

using clk = std::chrono::steady_clock;

static std::vector<float> random_vector(int dim, std::mt19937 &rng) {
  std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
  std::vector<float> v(dim);
  for (int i = 0; i < dim; i++)
    v[i] = dist(rng);
  return v;
}

static float cosine(const std::vector<float> &a, const std::vector<float> &b) {
  float dot = 0, am = 0, bm = 0;
  for (size_t i = 0; i < a.size(); i++) {
    dot += a[i] * b[i];
    am += a[i] * a[i];
    bm += b[i] * b[i];
  }
  if (am == 0 || bm == 0)
    return 0.0f; // guard against zero-magnitude vectors
  return dot / (std::sqrt(am) * std::sqrt(bm));
}

// Brute-force top-k by cosine similarity — ground truth for recall measurement.
static std::vector<int>
brute_force_topk(const std::vector<std::vector<float>> &data,
                 const std::vector<float> &query, int k) {
  std::vector<std::pair<float, int>> scored;
  scored.reserve(data.size());
  for (size_t i = 0; i < data.size(); i++)
    scored.push_back({cosine(data[i], query), (int)i});
  std::sort(scored.begin(), scored.end(),
            [](const auto &a, const auto &b) { return a.first > b.first; });

  std::vector<int> ids;
  for (int i = 0; i < std::min((int)scored.size(), k); i++)
    ids.push_back(scored[i].second);
  return ids;
}

// Tells you WHERE the bytes are, which the allocator counters can't.
// Reads vec2embeds/vec2docs/vec2neighbors directly — if those are private,
// add `friend MemoryReport report(const Hnsw&);` to Hnsw.hpp.
struct MemoryReport {
  std::size_t nodes = 0, edges = 0, levels = 0;
  std::size_t payload_used = 0, payload_reserved = 0;
  std::size_t adjacency_used = 0, adjacency_reserved = 0;
  std::size_t control = 0, inline_structs = 0, labels = 0;
};
static MemoryReport report(const Hnsw &ix) {
  MemoryReport r;
  // capacity(), not size() — this is the 25% you're losing to vector doubling
  r.inline_structs +=
      ix.vec2embeds.capacity() * sizeof(Embedding) +
      ix.vec2docs.capacity() * sizeof(Document) +
      ix.vec2neighbors.capacity() * sizeof(std::vector<std::vector<int>>);
  for (const auto &e : ix.vec2embeds) {
    r.nodes++;
    r.payload_used += e.components.size() * sizeof(float);
    r.payload_reserved += e.components.capacity() * sizeof(float);
  }
  for (const auto &d : ix.vec2docs)                // SSO means this is 0 at
    if (d.document_title.capacity() > 15)          // your label lengths — the
      r.labels += d.document_title.capacity() + 1; // check earns its keep later
  for (const auto &node : ix.vec2neighbors) {
    r.levels += node.size();
    r.control +=
        node.capacity() * sizeof(std::vector<int>); // inner vector headers
    for (const auto &lvl : node) {
      r.edges += lvl.size();
      r.adjacency_used += lvl.size() * sizeof(int);
      r.adjacency_reserved += lvl.capacity() * sizeof(int);
    }
  }
  return r;
}

static void run_benchmark(int n, int dim, int max_connections, int num_queries,
                          int k) {
  std::mt19937 rng(42); // fixed seed so runs are comparable across code changes
  std::cout << "\n=== N=" << n << ", dim=" << dim
            << ", max_connections=" << max_connections << " ===\n";

  // Corpus is generated up front, before any measurement window opens: its
  // allocations would otherwise land inside the index's memory delta, and its
  // RNG cost was previously inside the insert timer.
  std::vector<std::vector<float>> corpus;
  corpus.reserve(n);
  for (int i = 0; i < n; i++)
    corpus.push_back(random_vector(dim, rng));
  const std::size_t raw_bytes = (std::size_t)n * dim * sizeof(float);

  Hnsw index(max_connections, dim);
#ifdef HNSW_HAS_RESERVE
  index.reserve(n); // inside the window below, so the saving is measured
#endif

  // ---- Insert benchmark ----
  memprobe::reset_peak(); // peak becomes relative to this region
  auto m0 = memprobe::snap();
  auto t0 = clk::now();
  for (int i = 0; i < n; i++) {
    index.insert(corpus[i], i, "doc-" + std::to_string(i));
  }
  auto t1 = clk::now();
  auto m1 = memprobe::snap();
  double insert_secs = std::chrono::duration<double>(t1 - t0).count();

  // live delta == the index and nothing else: build temporaries cancel out,
  // and the corpus was already resident before m0.
  const std::size_t index_bytes = m1.live - m0.live;
  const std::size_t build_peak = m1.peak - m0.live;
  const std::size_t build_churn = m1.total - m0.total;
  const std::size_t build_allocs = m1.allocs - m0.allocs;

  std::cout << "Insert:  " << std::fixed << std::setprecision(3) << insert_secs
            << "s total, " << std::setprecision(1) << (n / insert_secs)
            << " vectors/sec\n";
  // ---- Query benchmark + recall ----
  std::vector<double> latencies_ms;
  latencies_ms.reserve(num_queries);
  std::size_t q_bytes = 0, q_allocs = 0;
  int hits = 0, total = 0;
  for (int q = 0; q < num_queries; q++) {
    auto query_vec = random_vector(dim, rng);
    Embedding qe(-1, dim, query_vec);
    // Window covers index.query() only — the brute-force ground truth below
    // allocates heavily and would drown the signal.
    auto a0 = memprobe::snap();
    auto qt0 = clk::now();
    auto results = index.query(qe, k);
    auto qt1 = clk::now();
    auto a1 = memprobe::snap();
    // `total`, not `live`: an alloc+free pair still costs latency and page
    // faults.
    q_bytes += a1.total - a0.total;
    q_allocs += a1.allocs - a0.allocs;
    latencies_ms.push_back(
        std::chrono::duration<double, std::milli>(qt1 - qt0).count());
    auto truth = brute_force_topk(corpus, query_vec, k);
    std::vector<int> got;
    got.reserve(results.size());
    for (auto &r : results)
      got.push_back(r.candidate_id);
    for (int t : truth) {
      total++;
      if (std::find(got.begin(), got.end(), t) != got.end())
        hits++;
    }
  }
  std::sort(latencies_ms.begin(), latencies_ms.end());
  auto pct = [&](double p) -> double {
    if (latencies_ms.empty())
      return 0.0;
    // was (size_t)(p * size), which returns the (floor(p*n)+1)-th smallest
    size_t idx = std::min(latencies_ms.size() - 1,
                          (size_t)std::ceil(p * latencies_ms.size()) - 1);
    return latencies_ms[idx];
  };
  double avg = latencies_ms.empty() ? 0.0
                                    : std::accumulate(latencies_ms.begin(),
                                                      latencies_ms.end(), 0.0) /
                                          latencies_ms.size();
  std::cout << "Query:   avg=" << std::setprecision(3) << avg << "ms"
            << "  p50=" << pct(0.50) << "ms"
            << "  p95=" << pct(0.95) << "ms"
            << "  p99=" << pct(0.99) << "ms\n";
  std::cout << "Recall@" << k << ": " << std::setprecision(1)
            << (total ? 100.0 * hits / total : 0.0)
            << "%  (vs. brute-force ground truth)\n";

  // ---- Memory report ----
  // Two independent sources, because each misses what the other catches:
  // the allocator hook is exact but blind to arena overhead; RSS is ground
  // truth but coarse; the structural walk says where the bytes actually are.
  auto mr = report(index);
  const std::size_t accounted = mr.payload_reserved + mr.adjacency_reserved +
                                mr.control + mr.inline_structs + mr.labels;
  const std::size_t slack = (mr.payload_reserved - mr.payload_used) +
                            (mr.adjacency_reserved - mr.adjacency_used);

  std::cout << "Memory:  index=" << human((double)index_bytes) << "  "
            << std::setprecision(1) << (double)index_bytes / n << " B/vector  "
            << std::setprecision(2) << (double)index_bytes / raw_bytes
            << "x raw (" << human((double)raw_bytes) << ")\n";
  std::cout << "         payload=" << human((double)mr.payload_reserved)
            << "  graph=" << human((double)mr.adjacency_reserved)
            << "  ctrl=" << human((double)mr.control)
            << "  inline=" << human((double)mr.inline_structs)
            << "  labels=" << human((double)mr.labels)
            << "  slack=" << human((double)slack) << "\n";
  std::cout << "         levels/node=" << std::setprecision(2)
            << (double)mr.levels / mr.nodes
            << "  avg degree=" << (double)mr.edges / mr.nodes
            << "  B/edge=" << std::setprecision(1)
            << (mr.edges ? (double)mr.adjacency_reserved / mr.edges : 0.0)
            // signed, so an over-counting accounting bug shows up as negative
            << "  unaccounted=" << std::setprecision(1)
            << (index_bytes
                    ? 100.0 * ((double)index_bytes - (double)accounted) /
                          index_bytes
                    : 0.0)
            << "%\n";
  std::cout << "         build peak=" << human((double)build_peak) << " ("
            << std::setprecision(2) << (double)build_peak / index_bytes
            << "x final)"
            << "  churn=" << human((double)build_churn) << " ("
            << std::setprecision(1) << (double)build_churn / index_bytes
            << "x final)"
            << "  allocs=" << (double)build_allocs / n << "/insert\n";
  std::cout << "         query transient="
            << human((double)q_bytes / num_queries) << "/query, "
            << std::setprecision(1) << (double)q_allocs / num_queries
            << " allocs/query\n";
  auto rs = memprobe::rss();
  std::cout << "Process: rss=" << human((double)rs.current)
            << "  peak rss=" << human((double)rs.peak) << "  (includes the "
            << human((double)raw_bytes) << " corpus)\n";
}

int main() {
  // Tune these to taste. Keep them fixed across runs so the numbers are
  // comparable as you optimize the implementation.
  run_benchmark(/*n=*/1000, /*dim=*/64, /*max_connections=*/16,
                /*num_queries=*/200, /*k=*/10);
  run_benchmark(/*n=*/5000, /*dim=*/64, /*max_connections=*/16,
                /*num_queries=*/200, /*k=*/10);
  run_benchmark(/*n=*/20000, /*dim=*/64, /*max_connections=*/16,
                /*num_queries=*/100, /*k=*/10);
  return 0;
}