Implement jump ahead for XORWOW RNG (#1049)

* Implement jump ahead for XORWOW RNG

* Construct both RNG engines with params

* Update formatting and doc, fix Weyl sequence, fix params assignment

* Add tests and more jumps

* Address review feedback

app/demo-interactor/KNDemoKernel.cu

@@ -73,7 +73,7 @@ __global__ void move_kernel(DeviceCRef<ParamsData> const params,
     // Construct particle accessor from immutable and thread-local data
     ParticleTrackView particle(
         params.particle, states.particle, TrackSlotId(tid));
-    RngEngine rng(states.rng, TrackSlotId(tid));
+    RngEngine rng(params.rng, states.rng, TrackSlotId(tid));
 
     // Move to collision
     XsCalculator calc_xs(params.tables.xs, params.tables.reals);
@@ -109,7 +109,7 @@ __global__ void interact_kernel(DeviceCRef<ParamsData> const params,
     // Construct particle accessor from immutable and thread-local data
     ParticleTrackView particle(
         params.particle, states.particle, TrackSlotId(tid));
-    RngEngine rng(states.rng, TrackSlotId(tid));
+    RngEngine rng(params.rng, states.rng, TrackSlotId(tid));
 
     Detector detector(params.detector, states.detector);
 

app/demo-interactor/KNDemoKernel.hh

@@ -66,6 +66,7 @@ template<Ownership W, MemSpace M>
 struct ParamsData
 {
     ParticleParamsData<W, M> particle;
+    RngParamsData<W, M> rng;
     TableData<W, M> tables;
     KleinNishinaData kn_interactor;
     DetectorParamsData detector;
@@ -81,6 +82,7 @@ struct ParamsData
     {
         CELER_EXPECT(other);
         particle = other.particle;
+        rng = other.rng;
         tables = other.tables;
         kn_interactor = other.kn_interactor;
         return *this;

src/celeritas/CMakeLists.txt

@@ -67,6 +67,7 @@ list(APPEND SOURCES
   phys/Process.cc
   phys/ProcessBuilder.cc
   random/CuHipRngData.cc
+  random/CuHipRngParams.cc
   random/XorwowRngData.cc
   random/XorwowRngParams.cc
   track/SimParams.cc

src/celeritas/global/CoreTrackView.hh

@@ -213,7 +213,7 @@ CELER_FUNCTION auto CoreTrackView::make_physics_step_view() const
  */
 CELER_FUNCTION auto CoreTrackView::make_rng_engine() const -> RngEngine
 {
-    return RngEngine{states_.rng, this->track_slot_id()};
+    return RngEngine{params_.rng, states_.rng, this->track_slot_id()};
 }
 
 //---------------------------------------------------------------------------//

src/celeritas/random/CuHipRngData.cc

@@ -31,25 +31,30 @@ void resize(CuHipRngStateData<Ownership::value, M>* state,
     CELER_EXPECT(size > 0);
     CELER_EXPECT(M == MemSpace::host || celeritas::device());
 
-    using RngInit = CuHipRngInitializer;
-
     // Host-side RNG for creating seeds
     std::mt19937 host_rng(params.seed + stream.get());
     std::uniform_int_distribution<ull_int> sample_uniform_int;
 
     // Create seeds for device in host memory
-    StateCollection<RngInit, Ownership::value, MemSpace::host> host_seeds;
+    StateCollection<ull_int, Ownership::value, MemSpace::host> host_seeds;
     resize(&host_seeds, size);
-    for (RngInit& init : host_seeds[AllItems<RngInit>{}])
+    for (auto& seed : host_seeds[AllItems<ull_int>{}])
     {
-        init.seed = sample_uniform_int(host_rng);
+        seed = sample_uniform_int(host_rng);
     }
 
+    // Set up params on device to initialize the engine
+    HostVal<CuHipRngParamsData> host_data;
+    host_data.seed = params.seed;
+    CuHipRngParamsData<Ownership::value, M> data;
+    data = host_data;
+
     // Resize state data and assign
     resize(&state->rng, size);
     detail::CuHipRngInitData<Ownership::value, M> init_data;
     init_data.seeds = host_seeds;
-    detail::rng_state_init(make_ref(*state), make_const_ref(init_data));
+    detail::rng_state_init(
+        make_const_ref(data), make_ref(*state), make_const_ref(init_data));
 }
 
 //---------------------------------------------------------------------------//

src/celeritas/random/CuHipRngData.hh

@@ -59,28 +59,9 @@ using CuHipRngThreadState = CELER_RNG_PREFIX(randState_t);
 //---------------------------------------------------------------------------//
 /*!
  * Properties of the global random number generator.
- *
- * There is no persistent data needed on device or at runtime: the params are
- * only used for construction.
  */
 template<Ownership W, MemSpace M>
-struct CuHipRngParamsData;
-
-template<Ownership W>
-struct CuHipRngParamsData<W, MemSpace::device>
-{
-    /* no data on device */
-
-    //! Assign from another set of data
-    template<Ownership W2, MemSpace M2>
-    CuHipRngParamsData& operator=(CuHipRngParamsData<W2, M2> const&)
-    {
-        return *this;
-    }
-};
-
-template<Ownership W>
-struct CuHipRngParamsData<W, MemSpace::host>
+struct CuHipRngParamsData
 {
     //// DATA ////
 
@@ -106,7 +87,9 @@ struct CuHipRngParamsData<W, MemSpace::host>
  */
 struct CuHipRngInitializer
 {
-    ull_int seed;
+    ull_int seed{0};
+    ull_int subsequence{0};
+    ull_int offset{0};
 };
 
 //---------------------------------------------------------------------------//

src/celeritas/random/CuHipRngEngine.hh

@@ -31,16 +31,18 @@ class CuHipRngEngine
     //! \name Type aliases
     using result_type = unsigned int;
     using Initializer_t = CuHipRngInitializer;
+    using ParamsRef = NativeCRef<CuHipRngParamsData>;
     using StateRef = NativeRef<CuHipRngStateData>;
     //!@}
 
   public:
     // Construct from state
-    inline CELER_FUNCTION
-    CuHipRngEngine(StateRef const& state, TrackSlotId const& id);
+    inline CELER_FUNCTION CuHipRngEngine(ParamsRef const& params,
+                                         StateRef const& state,
+                                         TrackSlotId tid);
 
     // Initialize state from seed
-    inline CELER_FUNCTION CuHipRngEngine& operator=(Initializer_t const& s);
+    inline CELER_FUNCTION CuHipRngEngine& operator=(Initializer_t const&);
 
     // Sample a random number
     inline CELER_FUNCTION result_type operator()();
@@ -97,10 +99,12 @@ class GenerateCanonical<CuHipRngEngine, double>
  * Construct from state.
  */
 CELER_FUNCTION
-CuHipRngEngine::CuHipRngEngine(StateRef const& state, TrackSlotId const& id)
+CuHipRngEngine::CuHipRngEngine(ParamsRef const&,
+                               StateRef const& state,
+                               TrackSlotId tid)
 {
-    CELER_EXPECT(id < state.rng.size());
-    state_ = &state.rng[id];
+    CELER_EXPECT(tid < state.rng.size());
+    state_ = &state.rng[tid];
 }
 
 //---------------------------------------------------------------------------//
@@ -109,7 +113,7 @@ CuHipRngEngine::CuHipRngEngine(StateRef const& state, TrackSlotId const& id)
  */
 CELER_FUNCTION CuHipRngEngine& CuHipRngEngine::operator=(Initializer_t const& s)
 {
-    CELER_RNG_PREFIX(rand_init)(s.seed, 0, 0, state_);
+    CELER_RNG_PREFIX(rand_init)(s.seed, s.subsequence, s.offset, state_);
     return *this;
 }
 

src/celeritas/random/CuHipRngParams.cc

@@ -0,0 +1,30 @@
+//----------------------------------*-C++-*----------------------------------//
+// Copyright 2023 UT-Battelle, LLC, and other Celeritas developers.
+// See the top-level COPYRIGHT file for details.
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file celeritas/random/CuHipRngParams.cc
+//---------------------------------------------------------------------------//
+#include "CuHipRngParams.hh"
+
+#include <utility>
+
+#include "corecel/Assert.hh"
+#include "celeritas/random/CuHipRngData.hh"
+
+namespace celeritas
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Construct with a seed.
+ */
+CuHipRngParams::CuHipRngParams(unsigned int seed)
+{
+    HostVal<CuHipRngParamsData> host_data;
+    host_data.seed = seed;
+    CELER_ASSERT(host_data);
+    data_ = CollectionMirror<CuHipRngParamsData>{std::move(host_data)};
+}
+
+//---------------------------------------------------------------------------//
+}  // namespace celeritas

src/celeritas/random/CuHipRngParams.hh

@@ -7,6 +7,7 @@
 //---------------------------------------------------------------------------//
 #pragma once
 
+#include "corecel/data/CollectionMirror.hh"
 #include "corecel/data/ParamsDataInterface.hh"
 
 #include "CuHipRngData.hh"
@@ -24,28 +25,18 @@ class CuHipRngParams final : public ParamsDataInterface<CuHipRngParamsData>
 {
   public:
     // Construct with seed
-    explicit inline CuHipRngParams(unsigned int seed);
+    explicit CuHipRngParams(unsigned int seed);
 
     //! Access RNG properties for constructing RNG state
-    HostRef const& host_ref() const final { return host_ref_; }
+    HostRef const& host_ref() const final { return data_.host(); }
 
     //! Access data on device
-    DeviceRef const& device_ref() const final { return device_ref_; }
+    DeviceRef const& device_ref() const final { return data_.device(); }
 
   private:
-    HostRef host_ref_;
-    DeviceRef device_ref_;
+    // Host/device storage and reference
+    CollectionMirror<CuHipRngParamsData> data_;
 };
 
 //---------------------------------------------------------------------------//
-// INLINE DEFINITIONS
-//---------------------------------------------------------------------------//
-/*!
- * Construct with seed.
- */
-CuHipRngParams::CuHipRngParams(unsigned int seed)
-{
-    host_ref_.seed = seed;
-}
-
 }  // namespace celeritas

src/celeritas/random/XorwowRngData.hh

@@ -18,19 +18,37 @@ namespace celeritas
 //---------------------------------------------------------------------------//
 /*!
  * Persistent data for XORWOW generator.
- *
- * If we want to add the "discard" operation or support initialization with a
- * subsequence or offset, we can add the precomputed XORWOW jump matrices here.
  */
 template<Ownership W, MemSpace M>
 struct XorwowRngParamsData
 {
+    //// TYPES ////
+
+    using uint_t = unsigned int;
+    using JumpPoly = Array<uint_t, 5>;
+    using ArrayJumpPoly = Array<JumpPoly, 32>;
+
+    //// DATA ////
+
     // TODO: 256-bit seed used to generate initial states for the RNGs
     // For now, just 4 bytes (same as our existing cuda/hip interface)
-    Array<unsigned int, 1> seed;
+    Array<uint_t, 1> seed;
+
+    // Jump polynomials
+    ArrayJumpPoly jump;
+    ArrayJumpPoly jump_subsequence;
 
     //// METHODS ////
 
+    static CELER_CONSTEXPR_FUNCTION size_type num_words()
+    {
+        return JumpPoly{}.size();
+    }
+    static CELER_CONSTEXPR_FUNCTION size_type num_bits()
+    {
+        return 8 * sizeof(uint_t);
+    }
+
     //! Whether the data is assigned
     explicit CELER_FUNCTION operator bool() const { return true; }
 
@@ -40,10 +58,23 @@ struct XorwowRngParamsData
     {
         CELER_EXPECT(other);
         seed = other.seed;
+        jump = other.jump;
+        jump_subsequence = other.jump_subsequence;
         return *this;
     }
 };
 
+//---------------------------------------------------------------------------//
+/*!
+ * Initialize an RNG.
+ */
+struct XorwowRngInitializer
+{
+    ull_int seed{0};
+    ull_int subsequence{0};
+    ull_int offset{0};
+};
+
 //---------------------------------------------------------------------------//
 //! Individual RNG state
 struct XorwowState