[libc] Add N Threads Benchmark Helper

[jameshu15869]

This PR adds a BENCHMARK_N_THREADS() helper to register benchmarks with a specific number of threads. This PR replaces the flags used originally to allow any amount of threads.

[llvmbot]

@llvm/pr-subscribers-libc

Author: None (jameshu15869)

Changes

This PR adds a BENCHMARK_N_THREADS() helper to register benchmarks with a specific number of threads. This PR replaces the flags used originally to allow any amount of threads.

---

Full diff: https://github.com/llvm/llvm-project/pull/99834.diff

2 Files Affected:

• (modified) libc/benchmarks/gpu/LibcGpuBenchmark.cpp (+1-4)
• (modified) libc/benchmarks/gpu/LibcGpuBenchmark.h (+14-12)

diff --git a/libc/benchmarks/gpu/LibcGpuBenchmark.cpp b/libc/benchmarks/gpu/LibcGpuBenchmark.cpp
index 59de18c20417d..ac139cc21b0b6 100644
--- a/libc/benchmarks/gpu/LibcGpuBenchmark.cpp
+++ b/libc/benchmarks/gpu/LibcGpuBenchmark.cpp
@@ -151,10 +151,7 @@ void Benchmark::run_benchmarks() {
       all_results.reset();
 
     gpu::sync_threads();
-    if (!b->flags ||
-        ((b->flags & BenchmarkFlags::SINGLE_THREADED) && id == 0) ||
-        ((b->flags & BenchmarkFlags::SINGLE_WAVE) &&
-         id < gpu::get_lane_size())) {
+    if (b->num_threads == static_cast<uint32_t>(-1) || id < b->num_threads) {
       auto current_result = b->run();
       all_results.update(current_result);
     }
diff --git a/libc/benchmarks/gpu/LibcGpuBenchmark.h b/libc/benchmarks/gpu/LibcGpuBenchmark.h
index c07fab9ccfbe3..f5cf4822f6fd3 100644
--- a/libc/benchmarks/gpu/LibcGpuBenchmark.h
+++ b/libc/benchmarks/gpu/LibcGpuBenchmark.h
@@ -74,8 +74,6 @@ struct BenchmarkResult {
   clock_t total_time = 0;
 };
 
-enum BenchmarkFlags { SINGLE_THREADED = 0x1, SINGLE_WAVE = 0x2 };
-
 BenchmarkResult benchmark(const BenchmarkOptions &options,
                           cpp::function<uint64_t(void)> wrapper_func);
 
@@ -83,12 +81,13 @@ class Benchmark {
   const cpp::function<uint64_t(void)> func;
   const cpp::string_view suite_name;
   const cpp::string_view test_name;
-  const uint8_t flags;
+  const uint32_t num_threads;
 
 public:
   Benchmark(cpp::function<uint64_t(void)> func, char const *suite_name,
-            char const *test_name, uint8_t flags)
-      : func(func), suite_name(suite_name), test_name(test_name), flags(flags) {
+            char const *test_name, uint32_t num_threads)
+      : func(func), suite_name(suite_name), test_name(test_name),
+        num_threads(num_threads) {
     add_benchmark(this);
   }
 
@@ -108,18 +107,21 @@ class Benchmark {
 } // namespace benchmarks
 } // namespace LIBC_NAMESPACE_DECL
 
+// Passing -1 indicates the benchmark should be run with as many threads as
+// allocated by the user in the benchmark's CMake.
 #define BENCHMARK(SuiteName, TestName, Func)                                   \
   LIBC_NAMESPACE::benchmarks::Benchmark SuiteName##_##TestName##_Instance(     \
-      Func, #SuiteName, #TestName, 0)
+      Func, #SuiteName, #TestName, -1)
 
-#define SINGLE_THREADED_BENCHMARK(SuiteName, TestName, Func)                   \
+#define BENCHMARK_N_THREADS(SuiteName, TestName, Func, NumThreads)             \
   LIBC_NAMESPACE::benchmarks::Benchmark SuiteName##_##TestName##_Instance(     \
-      Func, #SuiteName, #TestName,                                             \
-      LIBC_NAMESPACE::benchmarks::BenchmarkFlags::SINGLE_THREADED)
+      Func, #SuiteName, #TestName, NumThreads)
+
+#define SINGLE_THREADED_BENCHMARK(SuiteName, TestName, Func)                   \
+  BENCHMARK_N_THREADS(SuiteName, TestName, Func, 1)
 
 #define SINGLE_WAVE_BENCHMARK(SuiteName, TestName, Func)                       \
-  LIBC_NAMESPACE::benchmarks::Benchmark SuiteName##_##TestName##_Instance(     \
-      Func, #SuiteName, #TestName,                                             \
-      LIBC_NAMESPACE::benchmarks::BenchmarkFlags::SINGLE_WAVE)
+  BENCHMARK_N_THREADS(SuiteName, TestName, Func,                               \
+                      LIBC_NAMESPACE::gpu::get_lane_size())
 
 #endif

[jhuber6]

Thanks for the cleanup