Exclude the iteration count from the benchmark name

PiperOrigin-RevId: 728718349
Change-Id: I45986b511e0c7a63f7ddfee45100541aca61296c

Author: andreas-abel

fleetbench/BUILD

@@ -107,6 +107,7 @@ cc_library(
     deps = [
         ":dynamic_registrar",
         ":maybe_llvmlibc",
+        "//fleetbench/common",
         "@com_google_absl//absl/flags:parse",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",

fleetbench/benchmark_main.cc

@@ -24,6 +24,7 @@
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "benchmark/benchmark.h"
+#include "fleetbench/common/common.h"
 #include "fleetbench/dynamic_registrar.h"
 #include "numa.h"
 #include "tcmalloc/malloc_extension.h"
@@ -84,6 +85,15 @@ int main(int argc, char* argv[]) {
   }
   fleetbench::DynamicRegistrar::Get()->Run();
 
-  benchmark::RunSpecifiedBenchmarks();
+  fleetbench::FleetbenchReporter display_reporter(
+      benchmark::CreateDefaultDisplayReporter());
+  if (!fleetbench::FindInCommandLine("--benchmark_out=")) {
+    benchmark::RunSpecifiedBenchmarks(&display_reporter);
+  } else {
+    std::unique_ptr<benchmark::BenchmarkReporter> default_file_reporter =
+        fleetbench::CreateDefaultFileReporter();
+    fleetbench::FleetbenchReporter file_reporter(default_file_reporter.get());
+    benchmark::RunSpecifiedBenchmarks(&display_reporter, &file_reporter);
+  }
   return 0;
 }

fleetbench/common/common.cc

@@ -18,6 +18,7 @@
 #include <filesystem>  // NOLINT
 #include <fstream>
 #include <iostream>
+#include <memory>
 #include <optional>
 #include <random>
 #include <string>
@@ -227,4 +228,44 @@ int GetCacheSize(int cache_level, absl::string_view cache_type) {
                 "the cache sizes (in bytes) manually.";
 }
 
+std::unique_ptr<benchmark::BenchmarkReporter> CreateDefaultFileReporter() {
+  if (fleetbench::FindInCommandLine("--benchmark_out_format=console")) {
+    benchmark::ConsoleReporter::OutputOptions output_opts =
+        benchmark::ConsoleReporter::OutputOptions::OO_None;
+    if (fleetbench::FindInCommandLine("--benchmark_counters_tabular", true) ||
+        fleetbench::FindInCommandLine("--benchmark_counters_tabular=1") ||
+        fleetbench::FindInCommandLine("--benchmark_counters_tabular=true")) {
+      output_opts = benchmark::ConsoleReporter::OutputOptions::OO_Tabular;
+    }
+    return std::make_unique<benchmark::ConsoleReporter>(output_opts);
+  }
+  if (fleetbench::FindInCommandLine("--benchmark_out_format=csv")) {
+    return std::make_unique<benchmark::CSVReporter>();
+  }
+  return std::make_unique<benchmark::JSONReporter>();
+}
+
+bool FindInCommandLine(absl::string_view s, bool exact_match) {
+  std::ifstream f;
+  f.open("/proc/self/cmdline");
+  std::string command_line;
+  f >> command_line;
+  f.close();
+
+  std::vector<absl::string_view> argv = absl::StrSplit(command_line, '\0');
+  for (const auto& arg : argv) {
+    if (absl::StartsWith(arg, s)) {
+      if (!exact_match || arg == s) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+bool UseExplicitIterationCounts() {
+  return !FindInCommandLine("--benchmark_min_time") &&
+         !FindInCommandLine("--benchmark_list_tests");
+}
+
 }  // namespace fleetbench

fleetbench/common/common.h

@@ -15,6 +15,7 @@
 #define THIRD_PARTY_FLEETBENCH_COMMON_COMMON_H_
 
 #include <filesystem>  // NOLINT
+#include <memory>
 #include <optional>
 #include <random>
 #include <string>
@@ -23,6 +24,7 @@
 #include "absl/container/btree_map.h"
 #include "absl/flags/declare.h"
 #include "absl/strings/string_view.h"
+#include "benchmark/benchmark.h"
 
 // Exposed for testing only.
 ABSL_DECLARE_FLAG(bool, fixed_seed);
@@ -90,5 +92,59 @@ std::string GetFleetbenchRuntimePath(absl::string_view path);
 
 int GetCacheSize(int cache_level, absl::string_view cache_type = "");
 
+// A BenchmarkReporter adapter that removes the iteration count from the
+// benchmark name.
+// By default, the benchmark framework creates a benchmark name that contains
+// `/iterations:x` if the iteration count is specified by calling
+// `->Iterations(x)` on the benchmark object, but not if it is determined
+// automatically, or specified by the `--benchmark_min_time` flag.
+class FleetbenchReporter : public benchmark::BenchmarkReporter {
+ public:
+  explicit FleetbenchReporter(benchmark::BenchmarkReporter* reporter)
+      : reporter_(reporter) {}
+
+  void ReportRuns(
+      const std::vector<benchmark::BenchmarkReporter::Run>& runs) override {
+    std::vector<benchmark::BenchmarkReporter::Run> runs_copy(runs);
+    for (auto& run : runs_copy) {
+      run.run_name.iterations = "";
+    }
+    reporter_->ReportRuns(runs_copy);
+  }
+
+  void ReportRunsConfig(double min_time, bool has_explicit_iters,
+                        benchmark::IterationCount iters) override {
+    reporter_->ReportRunsConfig(min_time, has_explicit_iters, iters);
+  }
+
+  bool ReportContext(
+      const benchmark::BenchmarkReporter::Context& context) override {
+    reporter_->SetOutputStream(&this->GetOutputStream());
+    reporter_->SetErrorStream(&this->GetErrorStream());
+    return reporter_->ReportContext(context);
+  }
+
+  void Finalize() override { reporter_->Finalize(); }
+
+ private:
+  benchmark::BenchmarkReporter* const reporter_;
+};
+
+// Returns a file reporter that corresponds to the reporter that
+// benchmark::RunSpecifiedBenchmarks creates when the file_reporter parameter
+// is not set.
+std::unique_ptr<benchmark::BenchmarkReporter> CreateDefaultFileReporter();
+
+// Check if any command line argument starts with `s`.
+bool FindInCommandLine(absl::string_view s, bool exact_match = false);
+
+// Returns whether explicit iteration counts can be used when registering
+// benchmarks. This is the case if the `--benchmark_min_time` flag and the
+// `--benchmark_list_tests` flag are not set. `--benchmark_list_tests` is
+// excluded here because it does not use our custom `FleetbenchReporter` (see
+// `benchmark_main.cc`), and thus, the printed benchmark names would contain
+// `/iterations:x` when explicit iteration counts are used.
+bool UseExplicitIterationCounts();
+
 }  // namespace fleetbench
 #endif  // THIRD_PARTY_FLEETBENCH_COMMON_COMMON_H_