Class PhiloxRandom

tensorflow/core/lib/random/philox_random.h:103–253

A class that encapsulates all the states for a random number generator using the philox_4x32_10 algorithm. Each invocation returns a 128-bit random bits in the form of four uint32. There are multiple variants of this algorithm, we picked the 4x32_10 version that is most suited for our applications. Since this class is meant to be copied between CPU to GPU, it maintains a value semantics. For exam

// 1. PhiloxRandom is trivially copyable.
// 2. PhiloxRandom is compilable by gcc and nvcc.
class PhiloxRandom {
 public:
  using ResultType = Array<uint32, 4>;
  using ResultElementType = uint32;
  // The number of elements that will be returned.
  static const int kResultElementCount = 4;
  // Cost of generation of a single element (in cycles).
  static const int kElementCost = 10;
  // The type for the 64-bit key stored in the form of two 32-bit uint
  // that are used in the diffusion process.
  using Key = Array<uint32, 2>;

  PHILOX_DEVICE_INLINE
  PhiloxRandom() {}

  PHILOX_DEVICE_INLINE
  explicit PhiloxRandom(uint64 seed) {
    key_[0] = static_cast<uint32>(seed);
    key_[1] = static_cast<uint32>(seed >> 32);
  }

  PHILOX_DEVICE_INLINE
  explicit PhiloxRandom(uint64 seed_lo, uint64 seed_hi) {
    key_[0] = static_cast<uint32>(seed_lo);
    key_[1] = static_cast<uint32>(seed_lo >> 32);
    counter_[2] = static_cast<uint32>(seed_hi);
    counter_[3] = static_cast<uint32>(seed_hi >> 32);
  }

  PHILOX_DEVICE_INLINE
  PhiloxRandom(ResultType counter, Key key) : counter_(counter), key_(key) {}

  PHILOX_DEVICE_INLINE
  ResultType const& counter() const { return counter_; }

  PHILOX_DEVICE_INLINE
  Key const& key() const { return key_; }

  // Skip the specified number of samples of 128-bits in the current stream.
  PHILOX_DEVICE_INLINE
  void Skip(uint64 count) {
    const uint32 count_lo = static_cast<uint32>(count);
    uint32 count_hi = static_cast<uint32>(count >> 32);

    counter_[0] += count_lo;
    if (counter_[0] < count_lo) {
      ++count_hi;
    }

    counter_[1] += count_hi;
    if (counter_[1] < count_hi) {
      if (++counter_[2] == 0) {
        ++counter_[3];
      }
    }
  }

  // Returns a group of four random numbers using the underlying Philox