philox.h

// @HEADER
// *******************************************************************************
//                                OpenRAND                                       *
//   A Performance Portable, Reproducible Random Number Generation Library       *
//                                                                               *
// Copyright (c) 2023, Michigan State University                                 *
//                                                                               *
// Permission is hereby granted, free of charge, to any person obtaining a copy  *
// of this software and associated documentation files (the "Software"), to deal *
// in the Software without restriction, including without limitation the rights  *
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell     *
// copies of the Software, and to permit persons to whom the Software is         *
// furnished to do so, subject to the following conditions:                      *
//                                                                               *
// The above copyright notice and this permission notice shall be included in    *
// all copies or substantial portions of the Software.                           *
//                                                                               *
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR    *
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,      *
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE   *
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER        *
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, *
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE *
// SOFTWARE.                                                                     *
//********************************************************************************
// @HEADER
 
#ifndef OPENRAND_Philox_H_
#define OPENRAND_Philox_H_
 
#include <openrand/base_state.h>
 
#include <array>
#include <cstdint>
#include <iostream>
#include <limits>
 
#define PHILOX_W0 0x9E3779B9
#define PHILOX_W1 0xBB67AE85
#define PHILOX_M0 0xD2511F53
#define PHILOX_M1 0xCD9E8D57
 
namespace openrand {
 
class Philox : public BaseRNG<Philox> {
 public:
  OPENRAND_DEVICE Philox(uint64_t seed, uint32_t ctr,
                         uint32_t global_seed = openrand::DEFAULT_GLOBAL_SEED,
                         uint32_t ctr1 = 0x12345)
      : seed_hi((uint32_t)(seed >> 32)),
        seed_lo((uint32_t)(seed & 0xFFFFFFFF)),
        ctr0(ctr),
        ctr1(ctr1),
        ctr2(global_seed) {
  }
 
  template <typename T = uint32_t>
  OPENRAND_DEVICE T draw() {
    generate();
 
    static_assert(std::is_same_v<T, uint32_t> || std::is_same_v<T, uint64_t>);
    if constexpr (std::is_same_v<T, uint32_t>) return _out[0];
 
    // Not wrapping this block in else{} would lead to compiler warning
    else {
      uint64_t res = (static_cast<uint64_t>(_out[0]) << 32) |
                     static_cast<uint64_t>(_out[1]);
      return static_cast<uint64_t>(res);
    }
  }
 
  openrand::uint4 draw_int4() {
    generate();
    return openrand::uint4{_out[0], _out[1], _out[2], _out[3]};
  }
 
  openrand::float4 draw_float4() {
    generate();
    return openrand::float4{
        u01<float, uint32_t>(_out[0]), u01<float, uint32_t>(_out[1]),
        u01<float, uint32_t>(_out[2]), u01<float, uint32_t>(_out[3])};
  }
 
 private:
  OPENRAND_DEVICE void generate() {
    uint32_t key[2] = {seed_hi, seed_lo};
    _out[0] = ctr0;
    _out[1] = ctr1;
    _out[2] = ctr2;
    _out[3] = _ctr;
 
    for (int r = 0; r < 10; r++) {
      if (r > 0) {
        key[0] += PHILOX_W0;
        key[1] += PHILOX_W1;
      }
      round(key, _out);
    }
    _ctr++;
  }
 
  inline OPENRAND_DEVICE uint32_t mulhilo(uint32_t L, uint32_t R,
                                          uint32_t *hip) {
    uint64_t product = static_cast<uint64_t>(L) * static_cast<uint64_t>(R);
    *hip = static_cast<uint32_t>(product >> 32);
    return static_cast<uint32_t>(product);
  }
 
  inline OPENRAND_DEVICE void round(const uint32_t (&key)[2],
                                    uint32_t (&ctr)[4]) {
    uint32_t hi0;
    uint32_t hi1;
    uint32_t lo0 = mulhilo(PHILOX_M0, ctr[0], &hi0);
    uint32_t lo1 = mulhilo(PHILOX_M1, ctr[2], &hi1);
    ctr[0] = hi1 ^ ctr[1] ^ key[0];
    ctr[1] = lo1;
    ctr[2] = hi0 ^ ctr[3] ^ key[1];
    ctr[3] = lo0;
  }
 
  // User provided seed and counter, constant throughout
  // the lifetime of the rng object
  const uint32_t seed_hi, seed_lo;
  const uint32_t ctr0, ctr1, ctr2;
  uint32_t _out[4];
 
 public:
  // internal counter to keep track of numbers generated by this instance of rng
  uint32_t _ctr = 0;
};  // class Philox
 
}  // namespace openrand
 
#endif  // OPENRAND_Philox_H_