AMReX_Random.H

Go to the documentation of this file.

#ifndef AMREX_RAND_H
#define AMREX_RAND_H
#include <AMReX_Config.H>
 
#include <AMReX.H>
#include <AMReX_GpuQualifiers.H>
#include <AMReX_ParallelDescriptor.H>
#include <AMReX_RandomEngine.H>
#include <limits>
#include <cmath>
#include <cstdint>
 
namespace amrex
{
    namespace random_util {
 
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    Real one_to_zero (Real v) noexcept { return (v == Real(1)) ? Real(0) : v; }
 
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    Real zero_to_one (Real v) noexcept { return (v == Real(0)) ? Real(1) : v; }
 
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    Real clamp_below_one (Real v) noexcept
    {
        // The constexpr equivalent of std::nextafter(Real(1), Real(0)), which
        // is not available in device code.  That this really is the immediate
        // predecessor of one follows from the two assertions below: it is less
        // than one, and the midpoint between it and one already rounds to one,
        // so no Real lies in between.  Tests/Base/Random additionally checks it
        // against std::nextafter itself on the host.
        constexpr Real almost_one
            = Real(1) - Real(0.5)*std::numeric_limits<Real>::epsilon();
        static_assert(almost_one < Real(1));
        static_assert(Real(1) - Real(0.25)*std::numeric_limits<Real>::epsilon() == Real(1));
        // No generator produces NaN, but note this maps it to almost_one
        // rather than propagating it, since the comparison is false for NaN.
        return (v < almost_one) ? v : almost_one;
    }
 
    }
 
    Real Random ();
 
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    Real Random (RandomEngine const& random_engine)
    {
#if defined (__SYCL_DEVICE_ONLY__)
        // clamped for the same reason as the host path: the documented [0,1)
        // of the vendor generator is not a bound we want to depend on
        mkl::rng::device::uniform<Real> distr;
        return random_util::clamp_below_one(mkl::rng::device::generate(distr, *random_engine.engine));
#else
        // hiprand/curand draw from (0,1].  Relocating the single endpoint
        // that is out of range converts that to the documented [0,1) exactly,
        // with no arithmetic.  Subtracting from one would also convert the
        // interval, but not exactly: for a draw below half an ULP of one,
        // `1 - draw` rounds back up to 1.0, and the subtraction additionally
        // collapses every value below 0.5 onto a 2^-24 grid.
#ifdef BL_USE_FLOAT
        AMREX_IF_ON_DEVICE((
                AMREX_HIP_OR_CUDA(
                        return random_util::one_to_zero(hiprand_uniform(random_engine.rand_state)); ,
                        return random_util::one_to_zero(curand_uniform(random_engine.rand_state));
                )
        ))
#else
        AMREX_IF_ON_DEVICE((
                AMREX_HIP_OR_CUDA(
                        return random_util::one_to_zero(hiprand_uniform_double(random_engine.rand_state)); ,
                        return random_util::one_to_zero(curand_uniform_double(random_engine.rand_state));
                )
        ))
#endif
        AMREX_IF_ON_HOST((
                amrex::ignore_unused(random_engine);
                return Random();
        ))
#endif
    }
 
    Real RandomPositive ();
 
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    Real RandomPositive (RandomEngine const& random_engine)
    {
#if defined (__SYCL_DEVICE_ONLY__)
        // oneMKL draws from [0,1); relocating the zero endpoint converts that
        // to (0,1] with no arithmetic, and stays correct even if the
        // generator reaches its own upper bound, since 1.0 is in range here
        mkl::rng::device::uniform<Real> distr;
        return random_util::zero_to_one(mkl::rng::device::generate(distr, *random_engine.engine));
#else
        // (0,1] is the native interval of hiprand/curand, so unlike Random()
        // this needs no endpoint relocated and no clamp: the draw is
        // returned exactly as the generator produced it.
#ifdef BL_USE_FLOAT
        AMREX_IF_ON_DEVICE((
                AMREX_HIP_OR_CUDA(
                        return hiprand_uniform(random_engine.rand_state); ,
                        return curand_uniform(random_engine.rand_state);
                )
        ))
#else
        AMREX_IF_ON_DEVICE((
                AMREX_HIP_OR_CUDA(
                        return hiprand_uniform_double(random_engine.rand_state); ,
                        return curand_uniform_double(random_engine.rand_state);
                )
        ))
#endif
        AMREX_IF_ON_HOST((
                amrex::ignore_unused(random_engine);
                return RandomPositive();
        ))
#endif
    }
 
    Real RandomNormal (Real mean, Real stddev);
 
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    Real RandomNormal (Real mean, Real stddev, RandomEngine const& random_engine)
    {
#if defined (__SYCL_DEVICE_ONLY__)
        mkl::rng::device::gaussian<Real> distr(mean, stddev);
        return mkl::rng::device::generate(distr, *random_engine.engine);
#else
#ifdef BL_USE_FLOAT
        AMREX_IF_ON_DEVICE((
                AMREX_HIP_OR_CUDA(
                        return stddev * hiprand_normal(random_engine.rand_state) + mean; ,
                        return stddev * curand_normal(random_engine.rand_state) + mean;
                )
        ))
#else
        AMREX_IF_ON_DEVICE((
                AMREX_HIP_OR_CUDA(
                        return stddev * hiprand_normal_double(random_engine.rand_state) + mean; ,
                        return stddev * curand_normal_double(random_engine.rand_state) + mean;
                )
        ))
#endif
        AMREX_IF_ON_HOST((
                amrex::ignore_unused(random_engine);
                return RandomNormal(mean, stddev);
        ))
#endif
    }
 
    unsigned int RandomPoisson (Real lambda);
 
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    unsigned int RandomPoisson (Real lambda, RandomEngine const& random_engine)
    {
#if defined (__SYCL_DEVICE_ONLY__)
        mkl::rng::device::poisson<unsigned int> distr(lambda);
        return mkl::rng::device::generate(distr, *random_engine.engine);
#else
        AMREX_IF_ON_DEVICE((
                AMREX_HIP_OR_CUDA(
                        return hiprand_poisson(random_engine.rand_state, lambda); ,
                        return curand_poisson(random_engine.rand_state, lambda);
                )
        ))
        AMREX_IF_ON_HOST((
                amrex::ignore_unused(random_engine);
                return RandomPoisson(lambda);
        ))
#endif
    }
 
    namespace random_util {
 
    AMREX_GPU_DEVICE AMREX_FORCE_INLINE
    Real RandomGamma_alpha_ge_1 (Real alpha, Real beta, RandomEngine const& random_engine)
    {
        AMREX_ASSERT(alpha >= 1);
        AMREX_ASSERT(beta > 0);
 
        Real x, v, u;
        Real d = alpha - 1.0_rt / 3.0_rt;
        Real c = (1.0_rt / 3.0_rt) / std::sqrt(d);
 
        while (true) {
            do {
                x = amrex::RandomNormal(0.0_rt, 1.0_rt, random_engine);
                v = 1.0_rt + c * x;
            } while (v <= 0.0_rt);
 
            v = v * v * v;
            u = amrex::RandomPositive(random_engine); // std::log(u) below
 
            if (u < 1.0_rt - 0.0331_rt * x * x * x * x) {
                break;
            }
 
            if (std::log(u) < 0.5_rt * x * x + d * (1.0_rt - v + std::log(v))) {
                break;
            }
        }
        return beta * d * v;
    }
    }
 
    Real RandomGamma (Real alpha, Real beta);
 
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    Real RandomGamma (Real alpha, Real beta, RandomEngine const& random_engine)
    {
        AMREX_ASSERT(alpha > 0);
        AMREX_ASSERT(beta > 0);
 
        AMREX_IF_ON_DEVICE((
        if (alpha < 1)
        {
            Real u = amrex::RandomPositive(random_engine); // std::pow(u, 1/alpha) below
            return amrex::random_util::RandomGamma_alpha_ge_1(1.0_rt + alpha, beta, random_engine) * std::pow(u, 1.0_rt / alpha);
        } else {
            return amrex::random_util::RandomGamma_alpha_ge_1(alpha, beta, random_engine);
        }
        ))
 
        AMREX_IF_ON_HOST((
            amrex::ignore_unused(random_engine);
            return RandomGamma(alpha, beta);
        ))
    }
 
    unsigned int Random_int (unsigned int n); // [0,n-1]
 
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    unsigned int Random_int (unsigned int n, RandomEngine const& random_engine)
    {
    if (n == 0) { return 0;}
#if defined(__SYCL_DEVICE_ONLY__)
        mkl::rng::device::uniform<unsigned int> distr(0,n);
        return mkl::rng::device::generate(distr, *random_engine.engine);
#else
        AMREX_IF_ON_DEVICE((
                unsigned int rand;
                constexpr unsigned int RAND_M = 4294967295; // 2**32-1
                do {
                AMREX_HIP_OR_CUDA( rand = hiprand(random_engine.rand_state);,
                                rand =  curand(random_engine.rand_state) );
                } while (rand >= (RAND_M - RAND_M % n));
                return rand % n;
        ))
        AMREX_IF_ON_HOST((
                amrex::ignore_unused(random_engine);
                return Random_int(n);
        ))
#endif
    }
 
    ULong Random_long (ULong n); // [0,n-1]
 
    void FillRandom (Real* p, Long N);
 
    void FillRandomNormal (Real* p, Long N, Real mean, Real stddev);
 
    namespace detail {
        inline ULong DefaultGpuSeed () {
            return ParallelDescriptor::MyProc()*1234567ULL + 12345ULL;
        }
    }
 
    void InitRandom (ULong cpu_seed, int nprocs=ParallelDescriptor::NProcs(),
                     ULong gpu_seed = detail::DefaultGpuSeed());
 
    void ResetRandomSeed (ULong cpu_seed, ULong gpu_seed = detail::DefaultGpuSeed());
 
    void SaveRandomState (std::ostream& os);
 
    void RestoreRandomState (std::istream& is, int nthreads_old, int nstep_old);
 
    void UniqueRandomSubset (Vector<int> &uSet, int setSize, int poolSize,
                             bool printSet = false);
 
    void DeallocateRandomSeedDevArray ();
}
 
#endif