[libc] Add N Threads Benchmark Helper (#99834)

Summary:
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.

Test Plan: 

Reviewers: 

Subscribers: 

Tasks: 

Tags: 


Differential Revision: https://phabricator.intern.facebook.com/D60251455

Author: jameshu15869

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);
     }

libc/benchmarks/gpu/LibcGpuBenchmark.h

@@ -74,21 +74,20 @@ 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);
 
 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