Simplify litert_rng type selection, Remove redundant and confusing factory type.

LiteRT-PiperOrigin-RevId: 774993305

Author: LukeBoyer

litert/cc/litert_rng.h

@@ -116,14 +116,11 @@ class RandomDevice {
 
 // Abstract base class for generating data of a certain type from a given rng
 // device, e.g. populating tensors and the like.
-template <typename D, template <typename> typename Dist, typename DeviceBase>
+template <typename D, template <typename> typename Dist>
 class DataGenerator {
  public:
   using DataType = D;
   using Wide = WideType<D>;
-  using Device = RandomDevice<DeviceBase>;
-
-  virtual DataType operator()(Device& rng) = 0;
 
   // Bounds of distribution.
   DataType Max() const { return dist_.max(); }
@@ -134,14 +131,13 @@ class DataGenerator {
 };
 
 // A data generator that generates data within a given range.
-template <typename D, template <typename> typename Dist, typename DeviceBase>
-class RangedGenerator final : public DataGenerator<D, Dist, DeviceBase> {
+template <typename D, template <typename> typename Dist>
+class RangedGenerator final : public DataGenerator<D, Dist> {
  private:
-  using Base = DataGenerator<D, Dist, DeviceBase>;
+  using Base = DataGenerator<D, Dist>;
 
  public:
   using typename Base::DataType;
-  using typename Base::Device;
   using typename Base::Wide;
 
   RangedGenerator() = default;
@@ -155,30 +151,32 @@ class RangedGenerator final : public DataGenerator<D, Dist, DeviceBase> {
   RangedGenerator(RangedGenerator&&) = default;
   RangedGenerator& operator=(RangedGenerator&&) = default;
 
-  DataType operator()(Device& rng) override { return this->dist_(rng); }
+  template <typename Rng>
+  DataType operator()(Rng& rng) {
+    return this->dist_(rng);
+  }
 };
 
 // A rangeless float generator that casts random bits to the given float type.
 // This generally produces higher quality floats more repersentative of the
 // target distribution than a ranged generator. Particularly covers more values
 // around zero and infinities.
-template <typename D, template <typename> typename Dist, typename DeviceBase,
-          typename Enable = void>
-class ReinterpretGenerator final : public DataGenerator<D, Dist, DeviceBase> {};
+template <typename D, template <typename> typename Dist, typename Enable = void>
+class ReinterpretGenerator final : public DataGenerator<D, Dist> {};
 
-template <typename D, template <typename> typename Dist, typename DeviceBase>
-class ReinterpretGenerator<D, Dist, DeviceBase,
+template <typename D, template <typename> typename Dist>
+class ReinterpretGenerator<D, Dist,
                            std::enable_if_t<std::is_floating_point_v<D>>>
-    final : public DataGenerator<D, Dist, DeviceBase> {
+    final : public DataGenerator<D, Dist> {
  private:
-  using Base = DataGenerator<D, Dist, DeviceBase>;
+  using Base = DataGenerator<D, Dist>;
 
  public:
   using typename Base::DataType;
-  using typename Base::Device;
   using typename Base::Wide;
 
-  DataType operator()(Device& rng) override {
+  template <typename Rng>
+  DataType operator()(Rng& rng) {
     DataType res;
     auto bits = rng();
     memcpy(&res, &bits, sizeof(res));
@@ -195,72 +193,23 @@ class ReinterpretGenerator<D, Dist, DeviceBase,
   ReinterpretGenerator& operator=(ReinterpretGenerator&&) = default;
 };
 
-// Recommended distribution for data generators.
-template <typename T>
-using Uniform =
-    SelectT<std::is_floating_point<T>, std::uniform_real_distribution<T>,
-            std::is_integral<T>, std::uniform_int_distribution<T>>;
-
-// Recommended engine for data generators.
-using DefaultEngine = std::mt19937_64;
-
-// Factory for creating data generators from just a data type with recommended
-// defaults.
-template <typename D, template <typename> typename Distribution = Uniform,
-          typename Engine = DefaultEngine>
-class DataGenerators {
-  // Exotic types not yet supported (e.g. quant, complex, half-precision etc).
-  static_assert(std::is_floating_point_v<D> || std::is_integral_v<D>);
+// DEFAULTS FOR DATA GENERATORS ////////////////////////////////////////////////
 
- private:
-  using GeneratorBase = DataGenerator<D, Distribution, Engine>;
-
- public:
-  using Reinterpret = ReinterpretGenerator<D, Uniform, Engine>;
-  using Ranged = RangedGenerator<D, Uniform, Engine>;
-  using Dataype = GeneratorBase::DataType;
-  using Wide = GeneratorBase::Wide;
-  using RandomDevice = GeneratorBase::Device;
-
-  DataGenerators() = default;
-  DataGenerators(const DataGenerators&) = default;
-  DataGenerators& operator=(const DataGenerators&) = default;
-  DataGenerators(DataGenerators&&) = default;
-  DataGenerators& operator=(DataGenerators&&) = default;
-
-  // Create a ranged generator with the given limits.
-  static auto Generator(Wide min, Wide max) { return Ranged(min, max); }
-
-  // Create a rangeless generator. Floating point types will leverage the
-  // reinterpretation generator, which is recommended.
-  static auto Generator() {
-    if constexpr (std::is_floating_point_v<D>) {
-      return Reinterpret();
-    } else {
-      return Ranged();
-    }
-  }
+template <typename D>
+using DefaultGenerator =
+    SelectT<std::is_floating_point<D>,
+            ReinterpretGenerator<D, std::uniform_real_distribution>,
+            std::is_integral<D>,
+            RangedGenerator<D, std::uniform_int_distribution>>;
 
-  // Initialize a random device with the proper types to work with generators.
-  template <typename... Args>
-  static auto Device(Args&&... args) {
-    return RandomDevice(std::forward<Args>(args)...);
-  }
+template <typename D>
+using DefaultRangedGenerator =
+    SelectT<std::is_floating_point<D>,
+            RangedGenerator<D, std::uniform_real_distribution>,
+            std::is_integral<D>,
+            RangedGenerator<D, std::uniform_int_distribution>>;
 
-  // Convenience method(s) to create a generator and device in a pair.
-  static auto GeneratorAndDevice() {
-    return std::make_pair(Generator(), Device());
-  }
-  static auto GeneratorAndDevice(int seed) {
-    return std::make_pair(Generator(), Device(seed));
-  }
-  static auto GeneratorAndDevice(Wide min, Wide max) {
-    return std::make_pair(Generator(min, max), Device());
-  }
-  static auto GeneratorAndDevice(int seed, Wide min, Wide max) {
-    return std::make_pair(Generator(min, max), Device(seed));
-  }
-};
+using DefaultDevice = RandomDevice<std::mt19937>;
 
 }  // namespace litert
 

litert/cc/litert_rng_test.cc

@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <random>
+#include <type_traits>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -74,10 +75,14 @@ TEST(LitertRngTestWithCustomRng, NoSeed) {
               HasSubstr("DummyRng(seed=<default>,"));
 }
 
-using LiteRtRngTest = RngTest<>;
+using LiteRtRngTest = RngTest;
 
 TEST_F(LiteRtRngTest, Ints) {
-  auto [gen, device] = GeneratorAndDevice<int>();
+  auto device = TracedDevice();
+  auto gen = DefaultGenerator<int>();
+  static_assert(
+      std::is_same_v<decltype(gen),
+                     RangedGenerator<int, std::uniform_int_distribution>>);
   for (int i = 0; i < kTestIters; ++i) {
     const auto val = gen(device);
     ASSERT_LE(val, gen.Max());
@@ -88,7 +93,11 @@ TEST_F(LiteRtRngTest, Ints) {
 TEST_F(LiteRtRngTest, IntsWithRange) {
   static constexpr auto kMin = 10;
   static constexpr auto kMax = 20;
-  auto [gen, device] = GeneratorAndDevice<int>(kMin, kMax);
+  auto device = TracedDevice();
+  auto gen = DefaultGenerator<int>(kMin, kMax);
+  static_assert(
+      std::is_same_v<decltype(gen),
+                     RangedGenerator<int, std::uniform_int_distribution>>);
   EXPECT_EQ(gen.Max(), kMax);
   EXPECT_EQ(gen.Min(), kMin);
   for (int i = 0; i < kTestIters; ++i) {
@@ -99,9 +108,13 @@ TEST_F(LiteRtRngTest, IntsWithRange) {
 }
 
 TEST_F(LiteRtRngTest, FloatsWithRange) {
-  static constexpr auto kMin = 10;
-  static constexpr auto kMax = 20;
-  auto [gen, device] = GeneratorAndDevice<float>(kMin, kMax);
+  static constexpr auto kMin = 10.0f;
+  static constexpr auto kMax = 20.0f;
+  auto device = TracedDevice();
+  auto gen = DefaultRangedGenerator<float>(kMin, kMax);
+  static_assert(
+      std::is_same_v<decltype(gen),
+                     RangedGenerator<float, std::uniform_real_distribution>>);
   EXPECT_EQ(gen.Max(), kMax);
   EXPECT_EQ(gen.Min(), kMin);
   for (int i = 0; i < kTestIters; ++i) {
@@ -112,7 +125,11 @@ TEST_F(LiteRtRngTest, FloatsWithRange) {
 }
 
 TEST_F(LiteRtRngTest, ReinterpretFloat) {
-  auto [gen, device] = GeneratorAndDevice<float>();
+  auto device = TracedDevice();
+  auto gen = DefaultGenerator<float>();
+  static_assert(std::is_same_v<
+                decltype(gen),
+                ReinterpretGenerator<float, std::uniform_real_distribution>>);
   for (int i = 0; i < kTestIters; ++i) {
     const auto val = gen(device);
     ASSERT_FALSE(std::isnan(val));
@@ -123,7 +140,8 @@ TEST_F(LiteRtRngTest, ReinterpretFloat) {
 }
 
 TEST_F(LiteRtRngTest, TestWithFuzz) {
-  auto [gen, device] = GeneratorAndDevice<int>();
+  auto device = TracedDevice();
+  auto gen = DefaultGenerator<int>();
   for (auto _ :
        FuzzBlock(std::chrono::milliseconds(50), kTestIters, kTestIters)) {
     const auto val = gen(device);
@@ -133,5 +151,4 @@ TEST_F(LiteRtRngTest, TestWithFuzz) {
 }
 
 }  // namespace
-
 }  // namespace litert

litert/test/rng_fixture.h

@@ -32,13 +32,7 @@ namespace litert::testing {
 // set from gtest, which can be configured from command line for
 // reproducibility. It can also be used to set up repeated blocks of code for
 // fuzzing.
-template <template <typename> typename Distribution = Uniform,
-          typename Engine = DefaultEngine>
 class RngTest : public ::testing::Test {
- private:
-  template <typename T>
-  using Fact = DataGenerators<T, Distribution, Engine>;
-
  public:
   void TearDown() override {
     for (const auto& block : fuzz_blocks_) {
@@ -49,19 +43,9 @@ class RngTest : public ::testing::Test {
   }
 
  protected:
-  template <typename T, typename... Args>
-  auto Generator(Args&&... args) {
-    return Fact<T>::Generator(std::forward<Args>(args)...);
-  }
-  template <typename T>
-  auto Device() {
-    return TraceSeedInfo(Fact<T>::Device(CurrentSeed()));
-  }
-  template <typename T, typename... Args>
-  auto GeneratorAndDevice(Args&&... generator_args) {
-    return std::make_pair(
-        Fact<T>::Generator(std::forward<Args>(generator_args)...),
-        TraceSeedInfo(Fact<T>::Device(CurrentSeed())));
+  template <typename Device = DefaultDevice>
+  auto TracedDevice() {
+    return TraceSeedInfo(Device(CurrentSeed()));
   }
 
   template <typename... Args>