(ml5717) Started IncoherentRngCore impl + forward RNG to GPU

Author: juntyr

necsim-core/src/rng.rs

@@ -1,13 +1,72 @@
 use crate::intrinsics::{floor, ln};
 
+use core::convert::AsMut;
+use core::default::Default;
+use core::ptr::copy_nonoverlapping;
+
+#[allow(clippy::module_name_repetitions)]
+pub trait RngCore: Sized + Clone {
+    type Seed: AsMut<[u8]> + Default + Sized;
+
+    #[must_use]
+    fn from_seed(seed: Self::Seed) -> Self;
+
+    #[must_use]
+    fn seed_from_u64(mut state: u64) -> Self {
+        // Implementation from:
+        // https://docs.rs/rand/0.7.3/rand/trait.SeedableRng.html#method.seed_from_u64
+
+        // We use PCG32 to generate a u32 sequence, and copy to the seed
+        const MUL: u64 = 6_364_136_223_846_793_005_u64;
+        const INC: u64 = 11_634_580_027_462_260_723_u64;
+
+        let mut seed = Self::Seed::default();
+        for chunk in seed.as_mut().chunks_mut(4) {
+            // We advance the state first (to get away from the input value,
+            // in case it has low Hamming Weight).
+            state = state.wrapping_mul(MUL).wrapping_add(INC);
+
+            // Use PCG output function with to_le to generate x:
+            #[allow(clippy::cast_possible_truncation)]
+            let xorshifted = (((state >> 18) ^ state) >> 27) as u32;
+            #[allow(clippy::cast_possible_truncation)]
+            let rot = (state >> 59) as u32;
+            let x = xorshifted.rotate_right(rot).to_le();
+
+            unsafe {
+                let p = &x as *const u32 as *const u8;
+                copy_nonoverlapping(p, chunk.as_mut_ptr(), chunk.len());
+            }
+        }
+
+        Self::from_seed(seed)
+    }
+
+    #[must_use]
+    fn sample_u64(&mut self) -> u64;
+}
+
+pub trait IncoherentRngCore: RngCore {
+    type Prime: AsMut<[u8]> + Sized;
+
+    fn prime_with(&mut self, prime: Self::Prime);
+}
+
 #[allow(clippy::inline_always, clippy::inline_fn_without_body)]
 #[contract_trait]
-pub trait Core {
+pub trait Rng: RngCore {
+    #[must_use]
+    #[inline]
     #[debug_ensures(ret >= 0.0_f64 && ret <= 1.0_f64, "samples U(0.0, 1.0)")]
-    fn sample_uniform(&mut self) -> f64;
-}
+    fn sample_uniform(&mut self) -> f64 {
+        // http://prng.di.unimi.it -> Generating uniform doubles in the unit interval
+        #[allow(clippy::cast_precision_loss)]
+        ((self.sample_u64() >> 11) as f64)
+            * f64::from_bits(0x3CA0_0000_0000_0000_u64) //0x1.0p-53
+    }
 
-pub trait Rng: Core {
+    #[must_use]
+    #[inline]
     #[debug_ensures(ret < length, "samples U(0, length - 1)")]
     fn sample_index(&mut self, length: usize) -> usize {
         // attributes on expressions are experimental
@@ -21,12 +80,16 @@ pub trait Rng: Core {
         index
     }
 
+    #[must_use]
+    #[inline]
     #[debug_requires(lambda > 0.0, "lambda > 0.0")]
     #[debug_ensures(ret >= 0.0, "samples Exp(lambda)")]
     fn sample_exponential(&mut self, lambda: f64) -> f64 {
         -ln(self.sample_uniform()) / lambda
     }
 
+    #[must_use]
+    #[inline]
     #[debug_requires(
         probability >= 0.0_f64 && probability <= 1.0_f64,
         "0.0 <= probability <= 1.0"
@@ -36,4 +99,4 @@ pub trait Rng: Core {
     }
 }
 
-impl<T: Core> Rng for T {}
+impl<R: RngCore> Rng for R {}

necsim-cuda/kernel/src/lib.rs

@@ -44,35 +44,43 @@ use necsim_core::cogs::{
     IncoherentLineageStore, LineageReference,
 };
 use necsim_core::reporter::NullReporter;
+use necsim_core::rng::Rng;
 use necsim_core::simulation::Simulation;
-use necsim_impls_no_std::rng::wyrng::WyRng;
+use necsim_impls_no_std::rng::cuda::CudaRng;
 use rust_cuda::common::RustToCuda;
 use rust_cuda::device::BorrowFromRust;
 use rustacuda_core::DeviceCopy;
 
 #[no_mangle]
 /// # Safety
 /// This CUDA kernel is unsafe as it is called with raw pointers
-pub unsafe extern "ptx-kernel" fn simulate(c_void_ptr: *mut core::ffi::c_void) {
-    use necsim_impls_no_std::cogs::active_lineage_sampler::independent::IndependentActiveLineageSampler;
-    use necsim_impls_no_std::cogs::coalescence_sampler::independent::IndependentCoalescenceSampler;
-    use necsim_impls_no_std::cogs::dispersal_sampler::in_memory::packed_alias::InMemoryPackedAliasDispersalSampler;
-    use necsim_impls_no_std::cogs::event_sampler::independent::IndependentEventSampler;
-    use necsim_impls_no_std::cogs::habitat::in_memory::InMemoryHabitat;
-    use necsim_impls_no_std::cogs::lineage_reference::in_memory::InMemoryLineageReference;
-    use necsim_impls_no_std::cogs::lineage_store::incoherent::in_memory::IncoherentInMemoryLineageStore;
+pub unsafe extern "ptx-kernel" fn simulate(
+    simulation_c_ptr: *mut core::ffi::c_void,
+    cuda_rng_c_ptr: *mut core::ffi::c_void,
+    max_steps: usize,
+) {
+    use necsim_impls_no_std::cogs::active_lineage_sampler::independent::IndependentActiveLineageSampler as ActiveLineageSampler;
+    use necsim_impls_no_std::cogs::coalescence_sampler::independent::IndependentCoalescenceSampler as CoalescenceSampler;
+    use necsim_impls_no_std::cogs::dispersal_sampler::in_memory::packed_alias::InMemoryPackedAliasDispersalSampler as DispersalSampler;
+    use necsim_impls_no_std::cogs::event_sampler::independent::IndependentEventSampler as EventSampler;
+    use necsim_impls_no_std::cogs::habitat::in_memory::InMemoryHabitat as Habitat;
+    use necsim_impls_no_std::cogs::lineage_reference::in_memory::InMemoryLineageReference as LineageReference;
+    use necsim_impls_no_std::cogs::lineage_store::incoherent::in_memory::IncoherentInMemoryLineageStore as LineageStore;
+    use necsim_impls_no_std::rng::aes::AesRng as Rng;
 
     simulate_generic(
-        c_void_ptr
+        simulation_c_ptr
             as *mut <Simulation<
-                InMemoryHabitat,
-                InMemoryPackedAliasDispersalSampler,
-                InMemoryLineageReference,
-                IncoherentInMemoryLineageStore<_>,
-                IndependentCoalescenceSampler<_, _, _>,
-                IndependentEventSampler<_, _, _, _>,
-                IndependentActiveLineageSampler<_, _, _, _>,
+                Habitat,
+                DispersalSampler,
+                LineageReference,
+                LineageStore<_>,
+                CoalescenceSampler<_, _, _>,
+                EventSampler<_, _, _, _>,
+                ActiveLineageSampler<_, _, _, _>,
             > as RustToCuda>::CudaRepresentation,
+        cuda_rng_c_ptr as *mut <CudaRng<Rng> as RustToCuda>::CudaRepresentation,
+        max_steps,
     )
 }
 
@@ -84,21 +92,21 @@ unsafe fn simulate_generic<
     C: CoalescenceSampler<H, R, S> + RustToCuda,
     E: EventSampler<H, D, R, S, C> + RustToCuda,
     A: ActiveLineageSampler<H, D, R, S, C, E> + RustToCuda,
+    G: Rng,
 >(
     simulation_ptr: *mut <Simulation<H, D, R, S, C, E, A> as RustToCuda>::CudaRepresentation,
+    cuda_rng_ptr: *mut <CudaRng<G> as RustToCuda>::CudaRepresentation,
+    max_steps: usize,
 ) {
     Simulation::with_borrow_from_rust_mut(simulation_ptr, |simulation| {
-        let max_steps: usize = 100;
-        #[allow(clippy::cast_sign_loss)]
-        let rng_seed: u64 = utils::index() as u64;
-
-        let mut rng = WyRng::from_seed(rng_seed);
-        let mut reporter = NullReporter;
+        CudaRng::with_borrow_from_rust_mut(cuda_rng_ptr, |cuda_rng| {
+            let mut reporter = NullReporter;
 
-        //println!("{:#?}", simulation);
+            //println!("{:#?}", simulation);
 
-        let (time, steps) = simulation.simulate_incremental(max_steps, &mut rng, &mut reporter);
+            let (time, steps) = simulation.simulate_incremental(max_steps, cuda_rng, &mut reporter);
 
-        println!("time = {:?}, steps = {}", F64(time), steps);
+            println!("time = {:?}, steps = {}", F64(time), steps);
+        })
     })
 }

necsim-cuda/src/lib.rs

@@ -28,6 +28,7 @@ use necsim_impls_no_std::cogs::event_sampler::independent::IndependentEventSampl
 use necsim_impls_no_std::cogs::habitat::in_memory::InMemoryHabitat;
 use necsim_impls_no_std::cogs::lineage_reference::in_memory::InMemoryLineageReference;
 use necsim_impls_no_std::cogs::lineage_store::incoherent::in_memory::IncoherentInMemoryLineageStore;
+use necsim_impls_no_std::rng::cuda::CudaRng;
 use necsim_impls_std::cogs::dispersal_sampler::in_memory::InMemoryDispersalSampler;
 
 macro_rules! with_cuda {
@@ -79,7 +80,7 @@ impl CudaSimulation {
         dispersal: &Array2D<f64>,
         speciation_probability_per_generation: f64,
         sample_percentage: f64,
-        _rng: &mut impl Rng,
+        rng: &mut impl Rng,
         _reporter: &mut impl Reporter<InMemoryHabitat, InMemoryLineageReference>,
     ) -> Result<(f64, usize)> {
         let habitat = InMemoryHabitat::new(habitat.clone());
@@ -104,6 +105,8 @@ impl CudaSimulation {
             .active_lineage_sampler(active_lineage_sampler)
             .build();
 
+        let mut cuda_rng = CudaRng::from_cloned(rng);
+
         //let (time, steps) = simulation.simulate(rng, reporter);
 
         let module_data = CString::new(include_str!(env!("KERNEL_PTX_PATH"))).unwrap();
@@ -135,15 +138,19 @@ impl CudaSimulation {
             //println!("{:?}", CurrentContext::get_resource_limit(ResourceLimit::MaxL2FetchGranularity));
 
             if let Err(err) = simulation.lend_to_cuda_mut(|simulation_mut_ptr| {
-                // Launching kernels is unsafe since Rust can't enforce safety - think of kernel launches
-                // as a foreign-function call. In this case, it is - this kernel is written in CUDA C.
-                unsafe {
-                    launch!(module.simulate<<<1, 1, 0, stream>>>(
-                        simulation_mut_ptr
-                    ))?;
-                }
-
-                stream.synchronize()
+                cuda_rng.lend_to_cuda_mut(|cuda_rng_mut_ptr| {
+                    // Launching kernels is unsafe since Rust can't enforce safety - think of kernel launches
+                    // as a foreign-function call. In this case, it is - this kernel is written in CUDA C.
+                    unsafe {
+                        launch!(module.simulate<<<1, 1, 0, stream>>>(
+                            simulation_mut_ptr,
+                            cuda_rng_mut_ptr,
+                            1_000_usize // max steps on GPU
+                        ))?;
+                    }
+
+                    stream.synchronize()
+                })
             }) {
                 eprintln!("Running kernel failed with {:#?}!", err);
             }

necsim-impls-no-std/Cargo.toml

@@ -16,6 +16,7 @@ array2d = { path = "../array2d-no-std" }
 contracts = "0.6.0"
 float_next_after = { path = "../float-next-after-no-std" }
 nanorand = { version = "0.4.4", default-features = false, features = ["wyrand"] }
+aes-soft = "0.6.3"
 
 rust-cuda-derive = { path = "../rust-cuda/rust-cuda-derive", optional = true }
 rustacuda_core = { version = "0.1.2", optional = true }

necsim-impls-no-std/src/cogs/active_lineage_sampler/independent/sampler.rs

@@ -57,8 +57,16 @@ impl<
             .lineage_store
             .extract_lineage_from_its_location(chosen_lineage_reference.clone());
 
+        // TODO: As we are only doing geometric sampling for now, need to immediately increment discrete time step
+        // TODO: How do we choose the time step for now?
+        // TODO: Need to prime incoherent RNG here with location, discrete time step and substep 0
+
+        // TODO: Need to get time to next event in while loop with exponential (simplest option)
+
         let event_time = time + rng.sample_exponential(0.5_f64);
 
+        // TODO: Need to prime incoherent RNG here with location, discrete time step and substep 0
+
         let unique_event_time: f64 = if event_time > time {
             event_time
         } else {

necsim-impls-no-std/src/rng/aes.rs

@@ -0,0 +1,73 @@
+use aes_soft::cipher::generic_array::GenericArray;
+use aes_soft::cipher::{BlockCipher, NewBlockCipher};
+use aes_soft::Aes128;
+
+#[allow(clippy::module_name_repetitions)]
+#[derive(Clone)]
+pub struct AesRng {
+    cipher: Aes128,
+    state: [u8; 16],
+    cached: bool,
+}
+
+impl necsim_core::rng::RngCore for AesRng {
+    type Seed = [u8; 16];
+
+    #[must_use]
+    #[inline]
+    fn from_seed(seed: Self::Seed) -> Self {
+        Self {
+            cipher: Aes128::new(GenericArray::from_slice(&seed)),
+            state: [0_u8; 16],
+            cached: false,
+        }
+    }
+
+    #[must_use]
+    #[inline]
+    fn sample_u64(&mut self) -> u64 {
+        self.cached ^= true;
+
+        if self.cached {
+            // one more u64 will be cached
+            self.cipher
+                .encrypt_block(GenericArray::from_mut_slice(&mut self.state));
+
+            u64::from_le_bytes([
+                self.state[0],
+                self.state[1],
+                self.state[2],
+                self.state[3],
+                self.state[4],
+                self.state[5],
+                self.state[6],
+                self.state[7],
+            ])
+        } else {
+            // one more u64 was cached
+            let rand_u64 = u64::from_le_bytes([
+                self.state[8],
+                self.state[9],
+                self.state[10],
+                self.state[11],
+                self.state[12],
+                self.state[13],
+                self.state[14],
+                self.state[15],
+            ]);
+
+            self.state[9] = self.state[9].wrapping_add(1);
+
+            rand_u64
+        }
+    }
+}
+
+impl necsim_core::rng::IncoherentRngCore for AesRng {
+    type Prime = [u8; 16];
+
+    fn prime_with(&mut self, prime: Self::Prime) {
+        self.state = prime;
+        self.cached = false;
+    }
+}

necsim-impls-no-std/src/rng/cuda.rs

@@ -0,0 +1,29 @@
+use necsim_core::rng::RngCore;
+
+#[allow(clippy::module_name_repetitions)]
+#[derive(Clone, RustToCuda, LendToCuda)]
+pub struct CudaRng<R: RngCore>(R);
+
+impl<R: RngCore> CudaRng<R> {
+    #[must_use]
+    #[inline]
+    pub fn from_cloned(rng: &mut R) -> Self {
+        Self(rng.clone())
+    }
+}
+
+impl<R: RngCore> RngCore for CudaRng<R> {
+    type Seed = <R as RngCore>::Seed;
+
+    #[must_use]
+    #[inline]
+    fn from_seed(seed: Self::Seed) -> Self {
+        Self(R::from_seed(seed))
+    }
+
+    #[must_use]
+    #[inline]
+    fn sample_u64(&mut self) -> u64 {
+        self.0.sample_u64()
+    }
+}

necsim-impls-no-std/src/rng/mod.rs

@@ -1 +1,5 @@
-pub mod wyrng;
+pub mod aes;
+pub mod wy;
+
+#[cfg(feature = "cuda")]
+pub mod cuda;