Add an option to set the random number generator seed to a specific value.

PiperOrigin-RevId: 605068471
Change-Id: I5ecf6ffcf0179db7ed71d5b80e5d0e9993267795

Author: andreas-abel

fleetbench/common/BUILD

@@ -11,6 +11,7 @@ cc_library(
         "@bazel_tools//tools/cpp/runfiles",
         "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
         "@com_google_benchmark//:benchmark",
     ],

fleetbench/common/common.cc

@@ -26,6 +26,7 @@
 
 #include "tools/cpp/runfiles/runfiles.h"
 #include "absl/flags/flag.h"
+#include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/strings/match.h"
 #include "absl/strings/numbers.h"
@@ -37,7 +38,11 @@
 using bazel::tools::cpp::runfiles::Runfiles;
 
 ABSL_FLAG(bool, fixed_seed, true,
-          "Use a fixed seed for random number generation.");
+          "Use a fixed seed for random number generation. The seed can be "
+          "specified with the --seed flag; the default is 0. If "
+          "--fixed_seed=false, random seeds are used.");
+ABSL_FLAG(std::optional<int>, seed, {},
+          "Seed for random number generation. Requires --fixed_seed=true.");
 
 ABSL_FLAG(std::optional<int>, L1_data_size, {},
           "Size of the L1 data cache in bytes. Determined automatically if the "
@@ -54,13 +59,21 @@ namespace fleetbench {
 using CacheInfo = benchmark::CPUInfo::CacheInfo;
 
 Random& Random::instance() {
-  static auto* instance = new Random(absl::GetFlag(FLAGS_fixed_seed));
+  static auto* instance = [] {
+    if (absl::GetFlag(FLAGS_seed).has_value()) {
+      CHECK(absl::GetFlag(FLAGS_fixed_seed))
+          << "--seed requires --fixed_seed=true";
+    }
+    return new Random(absl::GetFlag(FLAGS_fixed_seed),
+                      absl::GetFlag(FLAGS_seed).value_or(0));
+  }();
   return *instance;
 }
 
-Random::Random(bool fixed_seed) : fixed_seed_(fixed_seed) {
+Random::Random(bool fixed_seed, int seed)
+    : fixed_seed_(fixed_seed), seed_(seed) {
   if (fixed_seed_) {
-    rng_.seed(0);
+    rng_.seed(seed_);
   } else {
     std::random_device rd;
     rng_.seed(rd());
@@ -69,7 +82,7 @@ Random::Random(bool fixed_seed) : fixed_seed_(fixed_seed) {
 
 void Random::Reset() {
   if (fixed_seed_) {
-    rng_.seed(0);
+    rng_.seed(seed_);
   }
 }
 

fleetbench/common/common.h

@@ -15,28 +15,35 @@
 #define THIRD_PARTY_FLEETBENCH_COMMON_COMMON_H_
 
 #include <filesystem>  // NOLINT
+#include <optional>
 #include <random>
 #include <string>
-#include <utility>
 #include <vector>
 
+#include "absl/flags/declare.h"
 #include "absl/strings/string_view.h"
 
+// Exposed for testing only.
+ABSL_DECLARE_FLAG(bool, fixed_seed);
+ABSL_DECLARE_FLAG(std::optional<int>, seed);
+
 namespace fleetbench {
 
 // Wrapper around a random number generator.
-// If --fixed_seed is true (the default), the seed will be 0 for better
-// reproducibility.
+// If --fixed_seed is true (the default), the seed will be set to the provided
+// value for better reproducibility. Otherwise, the seed is set to random
+// values.
 class Random {
  public:
   static Random& instance();
   std::default_random_engine& rng() { return rng_; }
   void Reset();
 
  private:
-  explicit Random(bool fixed_seed);
+  explicit Random(bool fixed_seed, int seed);
 
   const bool fixed_seed_;
+  const int seed_;
   std::default_random_engine rng_;
 };
 

fleetbench/common/common_test.cc

@@ -76,4 +76,24 @@ TEST(ReadCsvTest, Mixed) {
   EXPECT_EQ(result[1][1], "2");
 }
 
+TEST(RandomTest, Seed) {
+  absl::SetFlag(&FLAGS_fixed_seed, true);
+  absl::SetFlag(&FLAGS_seed, 1);
+
+  int random1 = GetRNG()();
+  Random::instance().Reset();
+  int random2 = GetRNG()();
+  EXPECT_EQ(random1, random2);
+}
+
+TEST(DeathTest, SeedFlags) {
+  ASSERT_DEATH(
+      {
+        absl::SetFlag(&FLAGS_fixed_seed, false);
+        absl::SetFlag(&FLAGS_seed, 1);
+        GetRNG()();
+      },
+      "--seed requires --fixed_seed=true");
+}
+
 }  // namespace fleetbench