From 860f2b9ad9d3602599aff49a41d0131d2a350e00 Mon Sep 17 00:00:00 2001
From: Pedro Gonnet <gonnet@google.com>
Date: Tue, 17 Dec 2024 13:45:24 -0800
Subject: [PATCH] Rely on `benchmark`'s flag parsing instead of implementing
 our own.

PiperOrigin-RevId: 707235530
---
 CMakeLists.txt            |  2 +-
 bench/models/BUILD        |  1 +
 bench/models/benchmark.cc | 41 ++++-----------------------------------
 3 files changed, 6 insertions(+), 38 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3d2175382a2..c94020d7e74 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1865,7 +1865,7 @@ IF(XNNPACK_BUILD_BENCHMARKS)
     TARGET_LINK_LIBRARIES(models PRIVATE XNNPACK)
 
     ADD_EXECUTABLE(bench-models bench/models/benchmark.cc)
-    TARGET_INCLUDE_DIRECTORIES(bench-models PRIVATE bench)
+    TARGET_INCLUDE_DIRECTORIES(bench-models PRIVATE bench ${GOOGLEBENCHMARK_SOURCE_DIR})
     TARGET_LINK_LIBRARIES(bench-models PRIVATE
       bench-utils
       benchmark::benchmark
diff --git a/bench/models/BUILD b/bench/models/BUILD
index fd5d9243bf0..0604e9ba0b8 100644
--- a/bench/models/BUILD
+++ b/bench/models/BUILD
@@ -32,6 +32,7 @@ xnnpack_cxx_library(
 xnnpack_benchmark(
     name = "benchmark",
     srcs = ["benchmark.cc"],
+    features = ["-layering_check"],
     tags = xnnpack_slow_benchmark_tags(),
     deps = [
         ":models",
diff --git a/bench/models/benchmark.cc b/bench/models/benchmark.cc
index 9ec21af4624..65e8190027c 100644
--- a/bench/models/benchmark.cc
+++ b/bench/models/benchmark.cc
@@ -5,12 +5,10 @@
 
 #include <benchmark/benchmark.h>
 
-#include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <functional>
-#include <iostream>
 #include <memory>
 #include <vector>
 
@@ -19,10 +17,11 @@
 #include "xnnpack.h"
 #include "xnnpack/allocator.h"
 #include "xnnpack/subgraph.h"
+#include "src/commandlineflags.h"
 #include "pthreadpool.h"
 
-int FLAGS_num_threads = 1;
-uint32_t FLAGS_xnn_runtime_flags = 0;
+BM_DEFINE_int32(num_threads, 1);
+BM_DEFINE_int32(xnn_runtime_flags, 0);
 
 struct ModelRuntime {
   std::unique_ptr<xnn_subgraph, decltype(&xnn_delete_subgraph)> model;
@@ -236,32 +235,6 @@ BENCHMARK(QD8Attention)
 
 BENCHMARK(QS8MobileNetV2)->Unit(benchmark::kMicrosecond)->UseRealTime();
 
-int ProcessArgs(int& argc, char**& argv) {
-  for (int i = 1; i < argc;) {
-    if (strncmp(argv[i], "--num_threads=", 14) == 0) {
-      FLAGS_num_threads = atoi(argv[i] + 14);
-      if (FLAGS_num_threads <= 0) {
-        std::cerr << "Invalid --num_threads: " << FLAGS_num_threads << "\n";
-        return 1;
-      }
-      std::copy(argv + i + 1, argv + argc, argv + i);
-      argc -= 1;
-    } else if (strncmp(argv[i], "--xnn_runtime_flags=", 20) == 0) {
-      const char* v = argv[i] + 20;
-      if (strlen(v) > 2 && strncmp(v, "0x", 2) == 0) {
-        FLAGS_xnn_runtime_flags = strtoul(v + 2, nullptr, 16);
-      } else {
-        FLAGS_xnn_runtime_flags = strtoul(v, nullptr, 10);
-      }
-      std::copy(argv + i + 1, argv + argc, argv + i);
-      argc -= 1;
-    } else {
-      ++i;
-    }
-  }
-  return 0;
-}
-
 #ifdef BENCHMARK_ARGS_BOTTLENECK
 // We are provided with a main that will call this function
 extern "C" {
@@ -270,12 +243,6 @@ int BenchmarkArgBottleneck(int& argc, char**& argv) {
 }
 }
 #else
-int main(int argc, char** argv) {
-  ::benchmark::Initialize(&argc, argv);
-  int status = ProcessArgs(argc, argv);
-  if (status != 0) return status;
-  if (::benchmark::ReportUnrecognizedArguments(argc, argv)) return 1;
-  ::benchmark::RunSpecifiedBenchmarks();
-}
+BENCHMARK_MAIN();
 #endif