[libcxx-commits] [libcxx] [libc++] Fix query benchmarks for associative containers (PR #216065)

[Louis Dionne via libcxx-commits]

https://github.com/ldionne created https://github.com/llvm/llvm-project/pull/216065

We were calling get_key and the random number generator in the hot loop. Instead, create a pool of keys to draw from in the initial setup phase, and then query these keys in a round-robin fashion.

>From 5241b359e1babaf0e39f94d698c06cc28ca9fb70 Mon Sep 17 00:00:00 2001
From: Louis Dionne <ldionne.2 at gmail.com>
Date: Thu, 13 Aug 2026 09:42:26 -0400
Subject: [PATCH] [libc++] Fix query benchmarks for associative containers

We were calling get_key and the random number generator in the hot loop.
Instead, create a pool of keys to draw from in the initial setup phase,
and then query these keys in a round-robin fashion.
---
 .../associative_container_benchmarks.h        | 26 +++++++++++++++----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/libcxx/test/benchmarks/containers/associative/associative_container_benchmarks.h b/libcxx/test/benchmarks/containers/associative/associative_container_benchmarks.h
index 34cd1ca90b647..d83c8db9c3ad5 100644
--- a/libcxx/test/benchmarks/containers/associative/associative_container_benchmarks.h
+++ b/libcxx/test/benchmarks/containers/associative/associative_container_benchmarks.h
@@ -706,12 +706,28 @@ void associative_container_benchmarks(std::string container) {
   /////////////////////////
   auto query_bench = [=](auto func) {
     return [=](auto& st) TEST_ALIGN_BENCHMARK {
-      const std::size_t size = st.range(0);
-      std::vector<Value> in  = make_value_types(generate_unique_keys(size));
-      Container c(in.begin(), in.end());
-
+      const std::size_t size    = st.range(0);
+      std::vector<Key> keys     = generate_unique_keys(size);
+      std::vector<Value> values = make_value_types(keys);
+      Container c(values.begin(), values.end());
+
+      // Create a pool of key indices to draw in the benchmark below. Make the
+      // pool large enough to defeat a branch predictor even when benchmarking
+      // small sizes.
+      const std::size_t N_DRAWS = std::max<std::size_t>(4096, keys.size());
+      std::vector<std::size_t> draws;
+      draws.reserve(N_DRAWS);
+      for (std::size_t i = 0; i != N_DRAWS; ++i) {
+        draws.push_back(getRandomEngine()() % keys.size());
+      }
+
+      std::size_t i = 0;
       for (auto _ : st) {
-        auto result = func(c, get_key(in[getRandomEngine()() % in.size()]));
+        Key const& key = keys[draws[i]];
+        auto result    = func(c, key);
+        ++i;
+        if (i == N_DRAWS)
+          i = 0;
         benchmark::DoNotOptimize(c);
         benchmark::DoNotOptimize(result);
         benchmark::ClobberMemory();