PacketRandomEngine.h

Go to the documentation of this file.

 
#ifndef EIGENRAND_PACKET_RANDOM_ENGINE_H
#define EIGENRAND_PACKET_RANDOM_ENGINE_H
 
#include <array>
#include <random>
#include <type_traits>
#include <EigenRand/MorePacketMath.h>
#include <fstream>
 
namespace Eigen
{
    namespace Rand
    {
        namespace detail
        {
            template<typename T>
            auto test_integral_result_type(int)->std::integral_constant<bool, std::is_integral<typename T::result_type>::value>;
 
            template<typename T>
            auto test_integral_result_type(...)->std::false_type;
 
            template<typename T>
            auto test_intpacket_result_type(int)->std::integral_constant<bool, internal::IsIntPacket<typename T::result_type>::value>;
 
            template<typename T>
            auto test_intpacket_result_type(...)->std::false_type;
        }
 
        template<typename Ty>
        struct IsScalarRandomEngine : decltype(detail::test_integral_result_type<Ty>(0))
        {
        };
 
        template<typename Ty>
        struct IsPacketRandomEngine : decltype(detail::test_intpacket_result_type<Ty>(0))
        {
        };
 
        enum class RandomEngineType
        {
            none, scalar, packet
        };
 
        template<typename Ty>
        struct GetRandomEngineType : std::integral_constant <
            RandomEngineType,
            IsPacketRandomEngine<Ty>::value ? RandomEngineType::packet :
            (IsScalarRandomEngine<Ty>::value ? RandomEngineType::scalar : RandomEngineType::none)
        >
        {
        };
 
        template<typename Ty, size_t length, size_t alignment = 64>
        class AlignedArray
        {
        public:
            AlignedArray()
            {
                allocate();
                for (size_t i = 0; i < length; ++i)
                {
                    new (&aligned[i]) Ty();
                }
            }
 
            AlignedArray(const AlignedArray& o)
            {
                allocate();
                for (size_t i = 0; i < length; ++i)
                {
                    aligned[i] = o[i];
                }
            }
 
            AlignedArray(AlignedArray&& o)
            {
                std::swap(memory, o.memory);
                std::swap(aligned, o.aligned);
            }
 
            AlignedArray& operator=(const AlignedArray& o)
            {
                for (size_t i = 0; i < length; ++i)
                {
                    aligned[i] = o[i];
                }
                return *this;
            }
 
            AlignedArray& operator=(AlignedArray&& o)
            {
                std::swap(memory, o.memory);
                std::swap(aligned, o.aligned);
                return *this;
            }
 
            ~AlignedArray()
            {
                deallocate();
            }
 
            Ty& operator[](size_t i)
            {
                return aligned[i];
            }
 
            const Ty& operator[](size_t i) const
            {
                return aligned[i];
            }
 
            size_t size() const
            {
                return length;
            }
 
            Ty* data()
            {
                return aligned;
            }
 
            const Ty* data() const
            {
                return aligned;
            }
 
        private:
            void allocate()
            {
                memory = std::malloc(sizeof(Ty) * length + alignment);
                aligned = (Ty*)(((size_t)memory + alignment) & ~(alignment - 1));
            }
 
            void deallocate()
            {
                if (memory)
                {
                    for (size_t i = 0; i < length; ++i)
                    {
                        aligned[i].~Ty();
                    }
                    std::free(memory);
                    memory = nullptr;
                    aligned = nullptr;
                }
            }
 
            void* memory = nullptr;
            Ty* aligned = nullptr;
        };
 
#ifndef EIGEN_DONT_VECTORIZE
        template<typename Packet,
            int _Nx, int _Mx,
            int _Rx, uint64_t _Px,
            int _Ux, uint64_t _Dx,
            int _Sx, uint64_t _Bx,
            int _Tx, uint64_t _Cx,
            int _Lx, uint64_t _Fx>
            class MersenneTwister
        {
        public:
            using result_type = Packet;
 
            static constexpr int word_size = 64;
            static constexpr int state_size = _Nx;
            static constexpr int shift_size = _Mx;
            static constexpr int mask_bits = _Rx;
            static constexpr uint64_t parameter_a = _Px;
            static constexpr int output_u = _Ux;
            static constexpr int output_s = _Sx;
            static constexpr uint64_t output_b = _Bx;
            static constexpr int output_t = _Tx;
            static constexpr uint64_t output_c = _Cx;
            static constexpr int output_l = _Lx;
 
            static constexpr uint64_t default_seed = 5489U;
 
            MersenneTwister(uint64_t x0 = default_seed)
            {
                using namespace Eigen::internal;
                std::array<uint64_t, unpacket_traits<Packet>::size / 2> seeds;
                for (uint64_t i = 0; i < seeds.size(); ++i)
                {
                    seeds[i] = x0 + i;
                }
                seed(ploadu<Packet>((int*)seeds.data()));
            }
 
            MersenneTwister(Packet x0)
            {
                seed(x0);
            }
 
            void seed(Packet x0)
            {
                using namespace Eigen::internal;
                Packet prev = state[0] = x0;
                for (int i = 1; i < _Nx; ++i)
                {
                    prev = state[i] = pmuluadd64(pxor(prev, psrl64(prev, word_size - 2)), _Fx, i);
                }
                stateIdx = _Nx;
            }
 
            uint64_t min() const
            {
                return 0;
            }
 
            uint64_t max() const
            {
                return _wMask;
            }
 
            result_type operator()()
            {
                if (stateIdx == _Nx)
                    refill_upper();
                else if (2 * _Nx <= stateIdx)
                    refill_lower();
 
                using namespace Eigen::internal;
 
                Packet res = state[stateIdx++];
                res = pxor(res, pand(psrl64(res, _Ux), pseti64<Packet>(_Dx)));
                res = pxor(res, pand(psll64(res, _Sx), pseti64<Packet>(_Bx)));
                res = pxor(res, pand(psll64(res, _Tx), pseti64<Packet>(_Cx)));
                res = pxor(res, psrl64(res, _Lx));
                return res;
            }
 
            void discard(unsigned long long num)
            {
                for (; 0 < num; --num)
                {
                    operator()();
                }
            }
 
            typename internal::HalfPacket<Packet>::type half()
            {
                if (valid)
                {
                    valid = false;
                    return cache;
                }
                typename internal::HalfPacket<Packet>::type a;
                internal::split_two(operator()(), a, cache);
                valid = true;
                return a;
            }
 
        protected:
 
            void refill_lower()
            {
                using namespace Eigen::internal;
 
                auto hmask = pseti64<Packet>(_hMask),
                    lmask = pseti64<Packet>(_lMask),
                    px = pseti64<Packet>(_Px),
                    one = pseti64<Packet>(1);
 
                int i;
                for (i = 0; i < _Nx - _Mx; ++i)
                {
                    Packet tmp = por(pand(state[i + _Nx], hmask),
                        pand(state[i + _Nx + 1], lmask));
 
                    state[i] = pxor(pxor(
                        psrl64(tmp, 1),
                        pand(pcmpeq64(pand(tmp, one), one), px)),
                        state[i + _Nx + _Mx]
                    );
                }
 
                for (; i < _Nx - 1; ++i)
                {
                    Packet tmp = por(pand(state[i + _Nx], hmask),
                        pand(state[i + _Nx + 1], lmask));
 
                    state[i] = pxor(pxor(
                        psrl64(tmp, 1),
                        pand(pcmpeq64(pand(tmp, one), one), px)),
                        state[i - _Nx + _Mx]
                    );
                }
 
                Packet tmp = por(pand(state[i + _Nx], hmask),
                    pand(state[0], lmask));
                state[i] = pxor(pxor(
                    psrl64(tmp, 1),
                    pand(pcmpeq64(pand(tmp, one), one), px)),
                    state[_Mx - 1]
                );
                stateIdx = 0;
            }
 
            void refill_upper()
            {
                using namespace Eigen::internal;
 
                auto hmask = pseti64<Packet>(_hMask),
                    lmask = pseti64<Packet>(_lMask),
                    px = pseti64<Packet>(_Px),
                    one = pseti64<Packet>(1);
 
                for (int i = _Nx; i < 2 * _Nx; ++i)
                {
                    Packet tmp = por(pand(state[i - _Nx], hmask),
                        pand(state[i - _Nx + 1], lmask));
 
                    state[i] = pxor(pxor(
                        psrl64(tmp, 1),
                        pand(pcmpeq64(pand(tmp, one), one), px)),
                        state[i - _Nx + _Mx]
                    );
                }
            }
 
            AlignedArray<Packet, _Nx * 2> state;
            size_t stateIdx = 0;
            typename internal::HalfPacket<Packet>::type cache;
            bool valid = false;
 
            static constexpr uint64_t _wMask = (uint64_t)-1;
            static constexpr uint64_t _hMask = (_wMask << _Rx) & _wMask;
            static constexpr uint64_t _lMask = ~_hMask & _wMask;
        };
 
        template<typename Packet>
        using Pmt19937_64 = MersenneTwister<Packet, 312, 156, 31,
            0xb5026f5aa96619e9, 29,
            0x5555555555555555, 17,
            0x71d67fffeda60000, 37,
            0xfff7eee000000000, 43, 6364136223846793005>;
#endif
 
        template<typename UIntType, typename BaseRng, int numU64>
        class ParallelRandomEngineAdaptor
        {
            static_assert(GetRandomEngineType<BaseRng>::value != RandomEngineType::none, "BaseRng must be a kind of Random Engine.");
        public:
            using result_type = UIntType;
 
            ParallelRandomEngineAdaptor(size_t seed = BaseRng::default_seed)
            {
                for (int i = 0; i < num_parallel; ++i)
                {
                    rngs[i].~BaseRng();
                    new (&rngs[i]) BaseRng{ seed + i * u64_stride };
                }
            }
 
            ParallelRandomEngineAdaptor(const BaseRng& o)
            {
                for (int i = 0; i < num_parallel; ++i)
                {
                    rngs[i].~BaseRng();
                    new (&rngs[i]) BaseRng{ o };
                }
            }
 
            ParallelRandomEngineAdaptor(const ParallelRandomEngineAdaptor&) = default;
            ParallelRandomEngineAdaptor(ParallelRandomEngineAdaptor&&) = default;
 
            static constexpr result_type min()
            {
                return std::numeric_limits<result_type>::min();
            }
 
            static constexpr result_type max()
            {
                return std::numeric_limits<result_type>::max();
            }
 
            result_type operator()()
            {
                if (cnt >= buf_size)
                {
                    refill_buffer();
                }
                return buf[cnt++];
            }
 
            float uniform_real()
            {
                if (fcnt >= fbuf_size)
                {
                    refill_fbuffer();
                }
                return fbuf[fcnt++];
            }
        private:
 
            void refill_buffer()
            {
                cnt = 0;
                for (size_t i = 0; i < num_parallel; ++i)
                {
                    reinterpret_cast<typename BaseRng::result_type&>(buf[i * result_type_stride]) = rngs[i]();
                }
            }
 
            void refill_fbuffer()
            {
                fcnt = 0;
                for (size_t i = 0; i < num_parallel; ++i)
                {
                    auto urf = internal::bit_to_ur_float(rngs[i]());
                    reinterpret_cast<decltype(urf)&>(fbuf[i * u64_stride * 2]) = urf;
                }
            }
 
            static constexpr int u64_stride = sizeof(typename BaseRng::result_type) / sizeof(uint64_t);
            static constexpr int result_type_stride = sizeof(typename BaseRng::result_type) / sizeof(result_type);
            static constexpr int num_parallel = numU64 / u64_stride;
            static constexpr int byte_size = sizeof(uint64_t) * numU64;
            static constexpr size_t buf_size = byte_size / sizeof(result_type);
            static constexpr size_t fbuf_size = byte_size / sizeof(float);
 
            std::array<BaseRng, num_parallel> rngs;
            AlignedArray<result_type, buf_size> buf;
            AlignedArray<float, fbuf_size> fbuf;
            size_t cnt = buf_size, fcnt = fbuf_size;
        };
 
        template<typename UIntType, typename BaseRng>
        using PacketRandomEngineAdaptor = ParallelRandomEngineAdaptor<UIntType, BaseRng,
            sizeof(typename BaseRng::result_type) / sizeof(uint64_t)>;
 
        template<typename BaseRng>
        class RandomEngineWrapper : public BaseRng
        {
        public:
            using BaseRng::BaseRng;
 
            RandomEngineWrapper(const BaseRng& o) : BaseRng{ o }
            {
            }
 
            RandomEngineWrapper(BaseRng&& o) : BaseRng{ o }
            {
            }
 
            RandomEngineWrapper(size_t seed) : BaseRng{ seed }
            {
            }
 
            RandomEngineWrapper() = default;
            RandomEngineWrapper(const RandomEngineWrapper&) = default;
            RandomEngineWrapper(RandomEngineWrapper&&) = default;
 
            float uniform_real()
            {
                internal::bit_scalar<float> bs;
                return bs.to_ur(this->operator()());
            }
        };
 
        template<typename UIntType, typename Rng>
        using UniversalRandomEngine = typename std::conditional<
            IsPacketRandomEngine<typename std::remove_reference<Rng>::type>::value,
            PacketRandomEngineAdaptor<UIntType, typename std::remove_reference<Rng>::type>,
            typename std::conditional<
            IsScalarRandomEngine<typename std::remove_reference<Rng>::type>::value,
            RandomEngineWrapper<typename std::remove_reference<Rng>::type>,
            void
            >::type
        >::type;
 
        template<typename UIntType, typename Rng>
        UniversalRandomEngine<UIntType, Rng> makeUniversalRng(Rng&& rng)
        {
            return { std::forward<Rng>(rng) };
        }
 
#ifdef EIGEN_VECTORIZE_AVX2
        using Vmt19937_64 = Pmt19937_64<internal::Packet8i>;
#elif defined(EIGEN_VECTORIZE_AVX) || defined(EIGEN_VECTORIZE_SSE2)
        using Vmt19937_64 = Pmt19937_64<internal::Packet4i>;
#else
        using Vmt19937_64 = std::mt19937_64;
#endif
        template<typename UIntType = uint64_t>
        using P8_mt19937_64 = ParallelRandomEngineAdaptor<UIntType, Vmt19937_64, 8>;
    }
}
 
#endif