improve random

2026-04-10 11:34:33 +08:00 · 2024-01-31 08:16:29 +00:00
parent a9ddab3e06
commit d626762e3f
9 changed files with 405 additions and 151 deletions
--- a/Eigen/src/Core/MathFunctions.h
+++ b/Eigen/src/Core/MathFunctions.h
@@ -563,34 +563,6 @@ struct pow_impl<ScalarX, ScalarY, true> {
  }
 };

-/****************************************************************************
- * Implementation of random                                               *
- ****************************************************************************/
-
-template <typename Scalar, bool IsComplex, bool IsInteger>
-struct random_default_impl {};
-
-template <typename Scalar>
-struct random_impl : random_default_impl<Scalar, NumTraits<Scalar>::IsComplex, NumTraits<Scalar>::IsInteger> {};
-
-template <typename Scalar>
-struct random_retval {
-  typedef Scalar type;
-};
-
-template <typename Scalar>
-inline EIGEN_MATHFUNC_RETVAL(random, Scalar) random(const Scalar& x, const Scalar& y);
-template <typename Scalar>
-inline EIGEN_MATHFUNC_RETVAL(random, Scalar) random();
-
-template <typename Scalar>
-struct random_default_impl<Scalar, false, false> {
-  static inline Scalar run(const Scalar& x, const Scalar& y) {
-    return x + (y - x) * Scalar(std::rand()) / Scalar(RAND_MAX);
-  }
-  static inline Scalar run() { return run(Scalar(NumTraits<Scalar>::IsSigned ? -1 : 0), Scalar(1)); }
-};
-
 enum { meta_floor_log2_terminate, meta_floor_log2_move_up, meta_floor_log2_move_down, meta_floor_log2_bogus };

 template <unsigned int n, int lower, int upper>
@@ -769,56 +741,166 @@ struct count_bits_impl<

 #endif  // EIGEN_COMP_GNUC || EIGEN_COMP_CLANG

+template <typename BitsType>
+int log2_ceil(BitsType x) {
+  int n = CHAR_BIT * sizeof(BitsType) - clz(x);
+  bool powerOfTwo = (x & (x - 1)) == 0;
+  return x == 0 ? 0 : powerOfTwo ? n - 1 : n;
+}
+
+template <typename BitsType>
+int log2_floor(BitsType x) {
+  int n = CHAR_BIT * sizeof(BitsType) - clz(x);
+  return x == 0 ? 0 : n - 1;
+}
+
+/****************************************************************************
+ * Implementation of random                                               *
+ ****************************************************************************/
+
+// return a Scalar filled with numRandomBits beginning from the least significant bit
+template <typename Scalar>
+Scalar getRandomBits(int numRandomBits) {
+  using BitsType = typename numext::get_integer_by_size<sizeof(Scalar)>::unsigned_type;
+  enum : int {
+    StdRandBits = meta_floor_log2<(unsigned int)(RAND_MAX) + 1>::value,
+    ScalarBits = sizeof(Scalar) * CHAR_BIT
+  };
+  eigen_assert((numRandomBits >= 0) && (numRandomBits <= ScalarBits));
+  const BitsType mask = BitsType(-1) >> (ScalarBits - numRandomBits);
+  BitsType randomBits = BitsType(0);
+  for (int shift = 0; shift < numRandomBits; shift += StdRandBits) {
+    int r = std::rand();
+    randomBits |= static_cast<BitsType>(r) << shift;
+  }
+  // clear the excess bits
+  randomBits &= mask;
+  return numext::bit_cast<Scalar, BitsType>(randomBits);
+}
+
+template <typename Scalar, bool IsComplex, bool IsInteger>
+struct random_default_impl {};
+
+template <typename Scalar>
+struct random_impl : random_default_impl<Scalar, NumTraits<Scalar>::IsComplex, NumTraits<Scalar>::IsInteger> {};
+
+template <typename Scalar>
+struct random_retval {
+  typedef Scalar type;
+};
+
+template <typename Scalar>
+inline EIGEN_MATHFUNC_RETVAL(random, Scalar) random(const Scalar& x, const Scalar& y);
+template <typename Scalar>
+inline EIGEN_MATHFUNC_RETVAL(random, Scalar) random();
+
+template <typename Scalar>
+struct random_default_impl<Scalar, false, false> {
+  using BitsType = typename numext::get_integer_by_size<sizeof(Scalar)>::unsigned_type;
+  enum : int { MantissaBits = NumTraits<Scalar>::digits() - 1 };
+  static EIGEN_DEVICE_FUNC inline Scalar run(const Scalar& x, const Scalar& y, int numRandomBits = MantissaBits) {
+    Scalar half_x = Scalar(0.5) * x;
+    Scalar half_y = Scalar(0.5) * y;
+    Scalar result = (half_x + half_y) + (half_y - half_x) * run(numRandomBits);
+    // result is in the half-open interval [x, y) -- provided that x < y
+    return result;
+  }
+  static EIGEN_DEVICE_FUNC inline Scalar run(int numRandomBits = MantissaBits) {
+    eigen_assert(numRandomBits >= 0 && numRandomBits <= MantissaBits);
+    BitsType randomBits = getRandomBits<BitsType>(numRandomBits);
+    // if fewer than MantissaBits is requested, shift them to the left
+    randomBits <<= (MantissaBits - numRandomBits);
+    // randomBits is in the half-open interval [2,4)
+    randomBits |= numext::bit_cast<BitsType>(Scalar(2));
+    // result is in the half-open interval [-1,1)
+    Scalar result = numext::bit_cast<Scalar>(randomBits) - Scalar(3);
+    return result;
+  }
+};
+
+// TODO: fix this for PPC
+template <bool Specialize = sizeof(long double) == 2 * sizeof(uint64_t) && !EIGEN_ARCH_PPC>
+struct random_longdouble_impl {
+  enum : int {
+    Size = sizeof(long double),
+    MantissaBits = NumTraits<long double>::digits() - 1,
+    LowBits = MantissaBits > 64 ? 64 : MantissaBits,
+    HighBits = MantissaBits > 64 ? MantissaBits - 64 : 0
+  };
+  static EIGEN_DEVICE_FUNC inline long double run() {
+    EIGEN_USING_STD(memcpy)
+    uint64_t randomBits[2];
+    long double result = 2.0L;
+    memcpy(&randomBits, &result, Size);
+    randomBits[0] |= getRandomBits<uint64_t>(LowBits);
+    randomBits[1] |= getRandomBits<uint64_t>(HighBits);
+    memcpy(&result, &randomBits, Size);
+    result -= 3.0L;
+    return result;
+  }
+};
+template <>
+struct random_longdouble_impl<false> {
+  using Impl = random_impl<double>;
+  static EIGEN_DEVICE_FUNC inline long double run() { return static_cast<long double>(Impl::run()); }
+};
+
+template <>
+struct random_impl<long double> {
+  static EIGEN_DEVICE_FUNC inline long double run(const long double& x, const long double& y) {
+    long double half_x = 0.5L * x;
+    long double half_y = 0.5L * y;
+    long double result = (half_x + half_y) + (half_y - half_x) * run();
+    return result;
+  }
+  static EIGEN_DEVICE_FUNC inline long double run() { return random_longdouble_impl<>::run(); }
+};
+
 template <typename Scalar>
 struct random_default_impl<Scalar, false, true> {
-  static inline Scalar run(const Scalar& x, const Scalar& y) {
+  using BitsType = typename numext::get_integer_by_size<sizeof(Scalar)>::unsigned_type;
+  enum : int { ScalarBits = sizeof(Scalar) * CHAR_BIT };
+  static EIGEN_DEVICE_FUNC inline Scalar run(const Scalar& x, const Scalar& y) {
    if (y <= x) return x;
-    // ScalarU is the unsigned counterpart of Scalar, possibly Scalar itself.
-    typedef typename make_unsigned<Scalar>::type ScalarU;
-    // ScalarX is the widest of ScalarU and unsigned int.
-    // We'll deal only with ScalarX and unsigned int below thus avoiding signed
-    // types and arithmetic and signed overflows (which are undefined behavior).
-    typedef std::conditional_t<(ScalarU(-1) > unsigned(-1)), ScalarU, unsigned> ScalarX;
-    // The following difference doesn't overflow, provided our integer types are two's
-    // complement and have the same number of padding bits in signed and unsigned variants.
-    // This is the case in most modern implementations of C++.
-    ScalarX range = ScalarX(y) - ScalarX(x);
-    ScalarX offset = 0;
-    ScalarX divisor = 1;
-    ScalarX multiplier = 1;
-    const unsigned rand_max = RAND_MAX;
-    if (range <= rand_max)
-      divisor = (rand_max + 1) / (range + 1);
-    else
-      multiplier = 1 + range / (rand_max + 1);
-    // Rejection sampling.
+    const BitsType range = static_cast<BitsType>(y) - static_cast<BitsType>(x) + 1;
+    // handle edge case where [x,y] spans the entire range of Scalar
+    if (range == 0) return getRandomBits<Scalar>(ScalarBits);
+    // calculate the number of random bits needed to fill range
+    const int numRandomBits = log2_ceil(range);
+    BitsType randomBits;
    do {
-      offset = (unsigned(std::rand()) * multiplier) / divisor;
-    } while (offset > range);
-    return Scalar(ScalarX(x) + offset);
+      randomBits = getRandomBits<BitsType>(numRandomBits);
+      // if the random draw is outside [0, range), try again (rejection sampling)
+      // in the worst-case scenario, the probability of rejection is: 1/2 - 1/2^numRandomBits < 50%
+    } while (randomBits >= range);
+    Scalar result = x + static_cast<Scalar>(randomBits);
+    return result;
  }

-  static inline Scalar run() {
+  static EIGEN_DEVICE_FUNC inline Scalar run() {
 #ifdef EIGEN_MAKING_DOCS
    return run(Scalar(NumTraits<Scalar>::IsSigned ? -10 : 0), Scalar(10));
 #else
-    enum {
-      rand_bits = meta_floor_log2<(unsigned int)(RAND_MAX) + 1>::value,
-      scalar_bits = sizeof(Scalar) * CHAR_BIT,
-      shift = plain_enum_max(0, int(rand_bits) - int(scalar_bits)),
-      offset = NumTraits<Scalar>::IsSigned ? (1 << (plain_enum_min(rand_bits, scalar_bits) - 1)) : 0
-    };
-    return Scalar((std::rand() >> shift) - offset);
+    return getRandomBits<Scalar>(ScalarBits);
 #endif
  }
 };

+template <>
+struct random_impl<bool> {
+  static EIGEN_DEVICE_FUNC inline bool run(const bool& x, const bool& y) {
+    if (y <= x) return x;
+    return run();
+  }
+  static EIGEN_DEVICE_FUNC inline bool run() { return getRandomBits<int>(1) ? true : false; }
+};
+
 template <typename Scalar>
 struct random_default_impl<Scalar, true, false> {
-  static inline Scalar run(const Scalar& x, const Scalar& y) {
+  static EIGEN_DEVICE_FUNC inline Scalar run(const Scalar& x, const Scalar& y) {
    return Scalar(random(x.real(), y.real()), random(x.imag(), y.imag()));
  }
-  static inline Scalar run() {
+  static EIGEN_DEVICE_FUNC inline Scalar run() {
    typedef typename NumTraits<Scalar>::Real RealScalar;
    return Scalar(random<RealScalar>(), random<RealScalar>());
  }
@@ -1863,13 +1945,6 @@ EIGEN_DEVICE_FUNC inline bool isApproxOrLessThan(
 ***  The special case of the  bool type ***
 ******************************************/

-template <>
-struct random_impl<bool> {
-  static inline bool run() { return random<int>(0, 1) == 0 ? false : true; }
-
-  static inline bool run(const bool& a, const bool& b) { return random<int>(a, b) == 0 ? false : true; }
-};
-
 template <>
 struct scalar_fuzzy_impl<bool> {
  typedef bool RealScalar;
--- a/Eigen/src/Core/NumTraits.h
+++ b/Eigen/src/Core/NumTraits.h
@@ -206,9 +206,7 @@ struct GenericNumTraits {

  EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR static inline T highest() { return (numext::numeric_limits<T>::max)(); }

-  EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR static inline T lowest() {
-    return IsInteger ? (numext::numeric_limits<T>::min)() : static_cast<T>(-(numext::numeric_limits<T>::max)());
-  }
+  EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR static inline T lowest() { return (numext::numeric_limits<T>::lowest)(); }

  EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR static inline T infinity() { return numext::numeric_limits<T>::infinity(); }

--- a/Eigen/src/Core/arch/Default/BFloat16.h
+++ b/Eigen/src/Core/arch/Default/BFloat16.h
@@ -677,16 +677,22 @@ EIGEN_ALWAYS_INLINE std::ostream& operator<<(std::ostream& os, const bfloat16& v
 namespace internal {

 template <>
-struct random_default_impl<bfloat16, false, false> {
-  static inline bfloat16 run(const bfloat16& x, const bfloat16& y) {
-    return x + (y - x) * bfloat16(float(std::rand()) / float(RAND_MAX));
-  }
-  static inline bfloat16 run() { return run(bfloat16(-1.f), bfloat16(1.f)); }
+struct is_arithmetic<bfloat16> {
+  enum { value = true };
 };

 template <>
-struct is_arithmetic<bfloat16> {
-  enum { value = true };
+struct random_impl<bfloat16> {
+  enum : int { MantissaBits = 7 };
+  using Impl = random_impl<float>;
+  static EIGEN_DEVICE_FUNC inline bfloat16 run(const bfloat16& x, const bfloat16& y) {
+    float result = Impl::run(x, y, MantissaBits);
+    return bfloat16(result);
+  }
+  static EIGEN_DEVICE_FUNC inline bfloat16 run() {
+    float result = Impl::run(MantissaBits);
+    return bfloat16(result);
+  }
 };

 }  // namespace internal
--- a/Eigen/src/Core/arch/Default/Half.h
+++ b/Eigen/src/Core/arch/Default/Half.h
@@ -762,16 +762,22 @@ EIGEN_ALWAYS_INLINE std::ostream& operator<<(std::ostream& os, const half& v) {
 namespace internal {

 template <>
-struct random_default_impl<half, false, false> {
-  static inline half run(const half& x, const half& y) {
-    return x + (y - x) * half(float(std::rand()) / float(RAND_MAX));
-  }
-  static inline half run() { return run(half(-1.f), half(1.f)); }
+struct is_arithmetic<half> {
+  enum { value = true };
 };

 template <>
-struct is_arithmetic<half> {
-  enum { value = true };
+struct random_impl<half> {
+  enum : int { MantissaBits = 10 };
+  using Impl = random_impl<float>;
+  static EIGEN_DEVICE_FUNC inline half run(const half& x, const half& y) {
+    float result = Impl::run(x, y, MantissaBits);
+    return half(result);
+  }
+  static EIGEN_DEVICE_FUNC inline half run() {
+    float result = Impl::run(MantissaBits);
+    return half(result);
+  }
 };

 }  // end namespace internal