From 1f98efb7e0a34bbe47899439a5fb44ad71bde417 Mon Sep 17 00:00:00 2001
From: jonmeow <jperkins@google.com>
Date: Tue, 15 Oct 2024 10:35:13 -0700
Subject: [PATCH] Add more compile benchmark stats

---
 toolchain/driver/BUILD                 |  1 +
 toolchain/driver/compile_benchmark.cpp | 35 ++++++++++++++++----------
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/toolchain/driver/BUILD b/toolchain/driver/BUILD
index ddfc7dfbec798..242b6e126b5e8 100644
--- a/toolchain/driver/BUILD
+++ b/toolchain/driver/BUILD
@@ -67,6 +67,7 @@ cc_binary(
         "//testing/base:global_exe_path",
         "//testing/base:source_gen_lib",
         "//toolchain/install:install_paths_test_helpers",
+        "//toolchain/testing:compile_helper",
         "@google_benchmark//:benchmark",
         "@llvm-project//llvm:Support",
     ],
diff --git a/toolchain/driver/compile_benchmark.cpp b/toolchain/driver/compile_benchmark.cpp
index a1ce8ed91d269..9a92bdc74dda4 100644
--- a/toolchain/driver/compile_benchmark.cpp
+++ b/toolchain/driver/compile_benchmark.cpp
@@ -10,6 +10,7 @@
 #include "testing/base/source_gen.h"
 #include "toolchain/driver/driver.h"
 #include "toolchain/install/install_paths_test_helpers.h"
+#include "toolchain/testing/compile_helper.h"
 
 namespace Carbon::Testing {
 namespace {
@@ -96,18 +97,30 @@ static auto BM_CompileAPIFileDenseDecls(benchmark::State& state) -> void {
   int num_files = ComputeFileCount(target_lines);
   llvm::OwningArrayRef<std::string> sources(num_files);
 
-  // Create a collection of random source files. Average the actual number of
-  // lines resulting so we can use that to compute the compilation speed as a
-  // line-rate counter.
-  double avg_lines = 0.0;
+  // Create a collection of random source files. Compute average statistics for
+  // counters for compilation speed.
+  CompileHelper compile_helper;
+  double total_bytes = 0.0;
+  double total_tokens = 0.0;
+  double total_lines = 0.0;
   for (std::string& source : sources) {
     source = bench.gen().GenAPIFileDenseDecls(target_lines,
                                               SourceGen::DenseDeclParams{});
-    avg_lines += llvm::count(source, '\n');
-  }
-  avg_lines /= sources.size();
-
-  // Setup the sources as files for compilation.
+    total_bytes += source.size();
+    total_tokens += compile_helper.GetTokenizedBuffer(source).size();
+    total_lines += llvm::count(source, '\n');
+  };
+  state.counters["Bytes"] =
+      benchmark::Counter(total_bytes / sources.size(),
+                         benchmark::Counter::kIsIterationInvariantRate);
+  state.counters["Tokens"] =
+      benchmark::Counter(total_tokens / sources.size(),
+                         benchmark::Counter::kIsIterationInvariantRate);
+  state.counters["Lines"] =
+      benchmark::Counter(total_lines / sources.size(),
+                         benchmark::Counter::kIsIterationInvariantRate);
+
+  // Set up the sources as files for compilation.
   llvm::OwningArrayRef<std::string> file_names = bench.SetUpFiles(sources);
   CARBON_CHECK(static_cast<int>(file_names.size()) == num_files);
 
@@ -132,10 +145,6 @@ static auto BM_CompileAPIFileDenseDecls(benchmark::State& state) -> void {
       i += static_cast<ssize_t>(success);
     }
   }
-
-  // Compute the line-rate of these compilations.
-  state.counters["Lines"] = benchmark::Counter(
-      avg_lines, benchmark::Counter::kIsIterationInvariantRate);
 }
 
 // Benchmark from 256-line test cases through 256k line test cases, and for each