celeritas/random/XorwowRngData.hh

0001 //----------------------------------*-C++-*----------------------------------//
0002 // Copyright 2022-2024 UT-Battelle, LLC, and other Celeritas developers.
0003 // See the top-level COPYRIGHT file for details.
0004 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
0005 //---------------------------------------------------------------------------//
0006 //! \file celeritas/random/XorwowRngData.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include <cstdint>
0011
0012 #include "corecel/Assert.hh"
0013 #include "corecel/Macros.hh"
0014 #include "corecel/Types.hh"
0015 #include "corecel/cont/Array.hh"
0016 #include "corecel/cont/Span.hh"
0017 #include "corecel/data/Collection.hh"
0018
0019 namespace celeritas
0020 {
0021 //---------------------------------------------------------------------------//
0022 //! 32-bit unsigned integer type for xorwow
0023 using XorwowUInt = std::uint32_t;
0024 //! Seed type used to generate initial states for the RNG
0025 using XorwowSeed = Array<XorwowUInt, 1>;
0026
0027 //---------------------------------------------------------------------------//
0028 /*!
0029  * Persistent data for XORWOW generator.
0030  */
0031 template<Ownership W, MemSpace M>
0032 struct XorwowRngParamsData
0033 {
0034     //// TYPES ////
0035
0036     using JumpPoly = Array<XorwowUInt, 5>;
0037     using ArrayJumpPoly = Array<JumpPoly, 32>;
0038
0039     //// DATA ////
0040
0041     //! \todo Use full 256-bit seed to generate initial states for the RNGs
0042     //! For now, just 4 bytes (same as our existing cuda/hip interface)
0043     XorwowSeed seed;
0044
0045     // Jump polynomials
0046     ArrayJumpPoly jump;
0047     ArrayJumpPoly jump_subsequence;
0048
0049     //// METHODS ////
0050
0051     static CELER_CONSTEXPR_FUNCTION size_type num_words()
0052     {
0053         return JumpPoly{}.size();
0054     }
0055     static CELER_CONSTEXPR_FUNCTION size_type num_bits()
0056     {
0057         return 8 * sizeof(XorwowUInt);
0058     }
0059
0060     //! Whether the data is assigned
0061     explicit CELER_FUNCTION operator bool() const { return true; }
0062
0063     //! Assign from another set of data
0064     template<Ownership W2, MemSpace M2>
0065     XorwowRngParamsData& operator=(XorwowRngParamsData<W2, M2> const& other)
0066     {
0067         CELER_EXPECT(other);
0068         seed = other.seed;
0069         jump = other.jump;
0070         jump_subsequence = other.jump_subsequence;
0071         return *this;
0072     }
0073 };
0074
0075 //---------------------------------------------------------------------------//
0076 /*!
0077  * Initialize an RNG.
0078  */
0079 struct XorwowRngInitializer
0080 {
0081     Array<unsigned int, 1> seed{0};
0082     ull_int subsequence{0};
0083     ull_int offset{0};
0084 };
0085
0086 //---------------------------------------------------------------------------//
0087 //! Individual RNG state
0088 struct XorwowState
0089 {
0090     Array<XorwowUInt, 5> xorstate;  //!< x, y, z, w, v
0091     XorwowUInt weylstate;  //!< d
0092 };
0093
0094 //---------------------------------------------------------------------------//
0095 /*!
0096  * XORWOW generator states for all threads.
0097  */
0098 template<Ownership W, MemSpace M>
0099 struct XorwowRngStateData
0100 {
0101     //// TYPES ////
0102
0103     template<class T>
0104     using Items = Collection<T, W, M>;
0105     template<class T>
0106     using StateItems = StateCollection<T, W, M>;
0107
0108     //// DATA ////
0109
0110     StateItems<XorwowState> state;  //!< Track state [track]
0111
0112     //// METHODS ////
0113
0114     //! True if assigned
0115     explicit CELER_FUNCTION operator bool() const { return !state.empty(); }
0116
0117     //! State size
0118     CELER_FUNCTION size_type size() const { return state.size(); }
0119
0120     //! Assign from another set of states
0121     template<Ownership W2, MemSpace M2>
0122     XorwowRngStateData& operator=(XorwowRngStateData<W2, M2>& other)
0123     {
0124         CELER_EXPECT(other);
0125         state = other.state;
0126         return *this;
0127     }
0128 };
0129
0130 //---------------------------------------------------------------------------//
0131 // Initialize XORWOW states with well-distributed random data
0132 void initialize_xorwow(Span<XorwowState> state,
0133                        XorwowSeed const& seed,
0134                        StreamId stream);
0135
0136 //---------------------------------------------------------------------------//
0137 // Resize and seed the RNG states
0138 template<MemSpace M>
0139 void resize(XorwowRngStateData<Ownership::value, M>* state,
0140             HostCRef<XorwowRngParamsData> const& params,
0141             StreamId stream,
0142             size_type size);
0143
0144 //---------------------------------------------------------------------------//
0145 }  // namespace celeritas