prng.hpp

#ifndef __MC_PRNG_HPP__
#define __MC_PRNG_HPP__
 
#include "Kokkos_Macros.hpp"
#include <Kokkos_Random.hpp>
#include <common/common.hpp>
#include <common/traits.hpp>
#include <cstdint>
#include <mc/alias.hpp>
namespace MC
{
 
  pool_type get_pool(std::size_t seed = 0);
 
  uint64_t master_seed();
  uint64_t next_splitmix64(uint64_t& x) noexcept;

 
  KOKKOS_INLINE_FUNCTION auto
  sample_random_variables(const pool_type& pool, auto&& functor)
  {
    // Be sure that generator is passed to functor as reference
    // As functor can be lambda usually [](auto gen){return gen.rand();};
    // If auto is passed as value, value generated is UB. Need auto& to avoid
    // any problem
    using gen_t = typename pool_type::generator_type;
    using Functor = decltype(functor);
    static_assert(std::is_invocable_v<Functor, gen_t&>
                      && !std::is_invocable_v<Functor, gen_t>,
                  "Functor must accept generator by reference only");
 
    auto gen = pool.get_state();
    const auto result = functor(gen);
    pool.free_state(gen);
 
    return result;
  }
 
#define SAMPLE_RANDOM_VARIABLES(_random_pool_, ...)                            \
  auto _generator_state_ = _random_pool_.get_state();                          \
  __VA_ARGS__;                                                                 \
  _random_pool_.free_state(_generator_state_);

  class KPRNG // TODO remove deprecated
  {
  public:
    using pool_type = MC::pool_type;
    using generator_type = pool_type::generator_type;
 
    KPRNG() = default;
 
    explicit KPRNG(uint64_t seed);

    template <class Archive>
    void
    save(Archive& ar) const
    {
      ar(this->seed);
    }

    template <class Archive>
    void
    load(Archive& ar)
    {
      uint64_t _seed = 0;
      ar(_seed);
      this->random_pool = get_pool(_seed);
    }
 
    template <FloatingPointType T>
    KOKKOS_INLINE_FUNCTION T
    uniform() const
    {
      return this->uniform<T>(0, 1);
    }
 
    template <FloatingPointType T>
    KOKKOS_INLINE_FUNCTION T
    uniform(T a, T b) const
    {
      return sample_random_variables(
          random_pool,
          [a, b](auto gen)
          {
            T x;
            if constexpr (std::is_same_v<T, float>)
            {
              x = gen.frand(a, b);
            }
            else if constexpr (std::is_same_v<T, double>)
            {
              x = gen.drand(a, b);
            }
            return x;
          });
    }
 
    [[deprecated]] [[nodiscard]] Kokkos::View<double*, ComputeSpace>
    double_uniform(size_t n_sample, double a = 0., double b = 1.) const;
 
    template <FloatingPointType T, size_t n_r>
    Kokkos::View<T[n_r], ComputeSpace>
    random_view() const
    {
      Kokkos::View<T[n_r], ComputeSpace> A("random");
      Kokkos::fill_random(A, random_pool, 0., 1.);
      return A;
    }
 
    template <size_t n_r>
    [[deprecated]] KOKKOS_INLINE_FUNCTION std::array<double, n_r>
    double_uniform() const
    {
      return generate_uniform_impl<pool_type, n_r>(
          random_pool, std::make_index_sequence<n_r>{});
    }
 
    [[nodiscard]] KOKKOS_INLINE_FUNCTION uint64_t
    uniform_u(uint64_t a, uint64_t b) const
    {
 
      return sample_random_variables(
          random_pool, [a, b](auto& gen) { return gen.urand64(a, b); });
    }
 
    pool_type random_pool;
 
    [[nodiscard]] auto
    get_seed() const
    {
      return seed;
    } // TODO export this to result file
 
  private:
    std::size_t seed{};
    template <typename random_pool_t, size_t n_r, size_t... I>
    KOKKOS_INLINE_FUNCTION std::array<double, n_r>
    generate_uniform_impl(random_pool_t pool,
                          std::index_sequence<I...> /*unused*/) const
    {
      // Constexpr loopunrolling to fill the array
      auto generator = pool.get_state();
      std::array<double, n_r> res
          = { { (static_cast<void>(I), generator.drand(0., 1.))... } };
      pool.free_state(generator);
      return res;
    }
  };
 
  template <class ExecutionSpace,
            class ViewType,
            class RandomPool,
            class IndexType = int64_t,
            const std::size_t CHUNK_SIZE>
  void
  fill_random(const ExecutionSpace& exec,
              ViewType a,
              RandomPool g,
              typename ViewType::const_value_type begin,
              typename ViewType::const_value_type end)
  {
    int64_t LDA = a.extent(0);
 
    if (LDA > 0)
    {
      Kokkos::parallel_for(
          "Kokkos::fill_random",
          Kokkos::RangePolicy<ExecutionSpace>(
              exec, 0, (LDA + (CHUNK_SIZE - 1)) / CHUNK_SIZE),
          Kokkos::Impl::fill_random_functor_begin_end<ViewType,
                                                      RandomPool,
                                                      CHUNK_SIZE,
                                                      ViewType::rank,
                                                      IndexType>(
              a, g, begin, end));
    }
  }
 
} // namespace MC
 
#endif //__MC_PRNG_HPP__