Allow detecting benchmark dry runs and reduce benchmark test costs (#8013)

We use the dry-run mode for running benchmarks in a test. This is
important, but doesn't need to use large batches of files. Doing so
would cause the compilation benchmark to become excessively slow,
especially on GitHub runners.

Sadly, Google Benchmark doesn't expose its dry-run mode, but we already
have our own `main` function that injects custom logic for our
benchmarks. Extend that to detect the dry run flags and environment
variables so that we can customize the benchmark behavior in that mode.
Then we use that detection to reduce the batch size.

Also cleans up the counters computation to both be a bit better
organized (I think), and to avoid keeping the tokenized buffers we build
to count tokens around during the timed running of the benchmark.

Also cleans up a bug in some clamping code spotted while fixing this.

Assisted-by: Antigravity with Gemini

---------

Co-authored-by: Dana Jansens <danakj@orodu.net>

testing/base/BUILD

@@ -37,6 +37,7 @@ cc_library(
     testonly = 1,
     srcs = ["benchmark_main.cpp"],
     deps = [
+        ":benchmark_helpers",
         ":global_exe_path",
         "//common:init_llvm",
         "@abseil-cpp//absl/flags:parse",
@@ -48,7 +49,9 @@ cc_library(
 cc_library(
     name = "benchmark_helpers",
     testonly = 1,
+    srcs = ["benchmark_helpers.cpp"],
     hdrs = ["benchmark_helpers.h"],
+    deps = ["//common:check"],
 )
 
 cc_library(

testing/base/benchmark_helpers.cpp

@@ -0,0 +1,27 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "testing/base/benchmark_helpers.h"
+
+#include <optional>
+
+#include "common/check.h"
+
+static constinit std::optional<bool> benchmark_dry_run = {};
+
+namespace Carbon::Testing {
+
+auto IsBenchmarkDryRun() -> bool {
+  CARBON_CHECK(benchmark_dry_run,
+               "queried `IsBenchmarkDryRun` before it was set");
+  return *benchmark_dry_run;
+}
+
+auto SetBenchmarkDryRun(bool is_dry_run) -> void {
+  CARBON_CHECK(!benchmark_dry_run,
+               "called `SetBenchmarkDryRun` more than once");
+  benchmark_dry_run = is_dry_run;
+}
+
+}  // namespace Carbon::Testing

testing/base/benchmark_helpers.h

@@ -65,6 +65,20 @@ template <typename T>
   DoNotOptimize(container.data());
 }
 
+// Returns true if the benchmark binary was invoked in dry-run mode
+// (`--benchmark_dry_run` or `BENCHMARK_DRY_RUN`).
+//
+// Google Benchmark's `--benchmark_dry_run` caps `state.max_iterations` to `1`,
+// so benchmarks that use `state.KeepRunningBatch(n)` or multi-input setup still
+// execute a full batch of `n` unless they also clamp `n` on dry runs.
+auto IsBenchmarkDryRun() -> bool;
+
+// Sets whether the benchmark binary is running in dry-run mode.
+//
+// This function must only be called once for an execution, and before any
+// callers to `IsBenchmarkDryRun`. Typically, it is called from within `main`.
+auto SetBenchmarkDryRun(bool is_dry_run) -> void;
+
 }  // namespace Carbon::Testing
 
 #endif  // CARBON_TESTING_BASE_BENCHMARK_HELPERS_H_

testing/base/benchmark_main.cpp

@@ -4,15 +4,55 @@
 
 #include <benchmark/benchmark.h>
 
+#include <cstdlib>
+
 #include "absl/flags/parse.h"
 #include "common/init_llvm.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringRef.h"
+#include "testing/base/benchmark_helpers.h"
 #include "testing/base/global_exe_path.h"
 
+// Returns true if `value` is truthy under Google Benchmark's boolean flag
+// rules: anything that isn't clearly _false_ is treated as truthy.
+static auto IsTruthyBenchmarkFlagValue(llvm::StringRef value) -> bool {
+  if (value.empty()) {
+    return true;
+  }
+  if (value.size() == 1) {
+    char c = value[0];
+    return llvm::isAlnum(c) && !llvm::StringRef("0fFnN").contains(c);
+  }
+  return !value.equals_insensitive("false") &&
+         !value.equals_insensitive("no") && !value.equals_insensitive("off");
+}
+
+// Detects `--benchmark_dry_run` (and `BENCHMARK_DRY_RUN` in the environment)
+// using the same rules as `benchmark::Initialize`, which consumes the flag
+// without exposing its value through Google Benchmark's public API.
+static auto DetectBenchmarkDryRun(llvm::ArrayRef<char*> args) -> bool {
+  bool dry_run = false;
+  if (const char* env = std::getenv("BENCHMARK_DRY_RUN")) {
+    dry_run = IsTruthyBenchmarkFlagValue(env);
+  }
+  for (llvm::StringRef arg : args) {
+    if (arg == "--benchmark_dry_run") {
+      dry_run = true;
+    } else if (arg.consume_front("--benchmark_dry_run=")) {
+      dry_run = IsTruthyBenchmarkFlagValue(arg);
+    }
+  }
+  return dry_run;
+}
+
 auto main(int orig_argc, char** orig_argv) -> int {
   // Do LLVM's initialization first, this will also transform UTF-16 to UTF-8.
   Carbon::InitLLVM init_llvm(orig_argc, orig_argv);
 
   Carbon::Testing::SetExePath(orig_argv[0]);
+  Carbon::Testing::SetBenchmarkDryRun(
+      DetectBenchmarkDryRun(llvm::ArrayRef(orig_argv, orig_argc).drop_front()));
 
   // Inject a flag to override the defaults for benchmarks. This can still be
   // disabled by user arguments.

toolchain/benchmarking/compile_benchmark.cpp

@@ -296,16 +296,27 @@ static auto ImplFileParams() -> SourceGen::DenseDeclParams {
 // For simplicity, we compute a number of files from the target line count as a
 // heuristic, clamped to the `[min_files, max_files]` range provided by the
 // compilation harness.
-static auto ComputeFileCount(int target_lines, int min_files,
+static auto ComputeFileCount(Lang lang, int target_lines, int min_files,
                              [[maybe_unused]] int max_files) -> int {
+  if (IsBenchmarkDryRun()) {
+    // Dry runs only smoke-test that the benchmark runs cleanly, so generate and
+    // compile just two files. We do more than one to exercise our stability
+    // checks.
+    return 2;
+  }
   int file_count = (1024 * 1024) / target_lines;
+  if (lang == Lang::Cpp) {
+    // Reduce the number of files with C++ to balance the longer compile times.
+    file_count /= 2;
+  }
 #ifndef NDEBUG
   // Use a smaller number of files in debug builds where compiles are slower,
   // capping at the release-mode minimum.
-  return std::max(1, std::min(min_files, file_count));
+  file_count = std::min(min_files, file_count);
 #else
-  return std::max(min_files, std::min(max_files, file_count));
+  file_count = std::min(max_files, file_count);
 #endif
+  return std::max(min_files, file_count);
 }
 
 template <Lang L, Phase P, Mode M = Mode::InProcess>
@@ -314,61 +325,46 @@ static auto RunApiFileBenchmark(benchmark::State& state,
     -> void {
   using Compiler = CompilerFor<L, M>;
   Compiler bench;
-  CompileHelper carbon_compile_helper;
 
   int target_lines = state.range(0);
   int num_files =
-      ComputeFileCount(target_lines, Compiler::MinFiles, Compiler::MaxFiles);
-  if constexpr (L == Lang::Cpp) {
-    // Reduce the number of files with C++ to balance the longer compile times.
-    num_files /= 2;
-  }
+      ComputeFileCount(L, target_lines, Compiler::MinFiles, Compiler::MaxFiles);
 
   llvm::SmallVector<std::string> sources;
   sources.reserve(num_files);
-
-  double total_bytes = 0.0;
-  double total_lines = 0.0;
-  double total_tokens = 0.0;
-
-  ssize_t first_bytes = 0;
-  ssize_t first_lines = 0;
-
   for (auto _ : llvm::seq(num_files)) {
     sources.push_back(bench.gen().GenApiFileDenseDecls(target_lines, params));
-    const auto& source = sources.back();
-    ssize_t bytes = source.size();
-    ssize_t lines = llvm::count(source, '\n');
-
-    total_bytes += bytes;
-    total_lines += lines;
-    if constexpr (L == Lang::Carbon) {
-      total_tokens += carbon_compile_helper.GetTokenizedBuffer(source).size();
-    }
-
-    // The generator varies each file's content but not its size, which keeps
-    // results comparable across files and runs.
-    if (sources.size() == 1) {
-      first_bytes = bytes;
-      first_lines = lines;
-      continue;
-    }
-
-    CARBON_CHECK(bytes == first_bytes,
-                 "Generated file {0} has {1} bytes but file 0 has {2}.",
-                 sources.size() - 1, bytes, first_bytes);
-    CARBON_CHECK(lines == first_lines,
-                 "Generated file {0} has {1} lines but file 0 has {2}.",
-                 sources.size() - 1, lines, first_lines);
   }
 
+  ssize_t bytes = sources[0].size();
+  ssize_t lines = llvm::count(sources[0], '\n');
+
+  for (llvm::StringRef source : llvm::ArrayRef(sources).drop_front()) {
+    // The generator varies each file's content but not its size, which keeps
+    // results comparable across files and runs.
+    ssize_t file_bytes = source.size();
+    ssize_t file_lines = llvm::count(source, '\n');
+
+    CARBON_CHECK(file_bytes == bytes,
+                 "Generated file has {0} bytes but expected {1}.", file_bytes,
+                 bytes);
+    CARBON_CHECK(file_lines == lines,
+                 "Generated file has {0} lines but expected {1}.", file_lines,
+                 lines);
+  }
   state.counters["Bytes"] =
-      benchmark::Counter(total_bytes / sources.size(),
-                         benchmark::Counter::kIsIterationInvariantRate);
+      benchmark::Counter(bytes, benchmark::Counter::kIsIterationInvariantRate);
   state.counters["Lines"] =
-      benchmark::Counter(total_lines / sources.size(),
-                         benchmark::Counter::kIsIterationInvariantRate);
+      benchmark::Counter(lines, benchmark::Counter::kIsIterationInvariantRate);
+
+  // TODO: The token count should be stable as well, so we should be able to
+  // switch it to the pattern above.
   if constexpr (L == Lang::Carbon) {
+    double total_tokens = 0.0;
+    CompileHelper carbon_compile_helper;
+    for (llvm::StringRef source : sources) {
+      total_tokens += carbon_compile_helper.GetTokenizedBuffer(source).size();
+    }
     state.counters["Tokens"] =
         benchmark::Counter(total_tokens / sources.size(),
                            benchmark::Counter::kIsIterationInvariantRate);