cugraph/
random.rs

1//! Random number generator state for cuGraph algorithms.
2//!
3//! Safe owner for `cugraph_rng_state_t` on the bound CUDA device.
4//!
5//! # Example
6//!
7//! ```no_run
8//! use cudf::Runtime;
9//! use cugraph::{ResourceHandle, random::RngState};
10//!
11//! let runtime = std::sync::Arc::new(Runtime::new().unwrap());
12//! let handle =
13//!     ResourceHandle::for_runtime(std::sync::Arc::clone(&runtime)).unwrap();
14//! let entered = handle.enter("RngState docs").unwrap();
15//! let rng = RngState::new(&entered, 42).unwrap();
16//! ```
17
18use std::{marker::PhantomData, ptr};
19
20use cugraph_sys::{cugraph_error_t, cugraph_rng_state_create, cugraph_rng_state_t};
21
22use crate::{
23    cuda::{DeviceResourceKind, drop_device_resource},
24    error::{Result, check_error, require_non_null},
25    handle::{EnteredContext, ResourceHandleLease},
26};
27
28/// RAII wrapper around `cugraph_rng_state_t`.
29///
30/// The RNG state is used by stochastic algorithms and is updated in-place
31/// with each call. Not thread-safe; must be used from the thread that
32/// created it.
33pub struct RngState {
34    raw: *mut cugraph_rng_state_t,
35    device: i32,
36    // Keeps a runtime-bound cuGraph handle alive until after Drop frees `raw`.
37    _resource_handle_lease: Option<ResourceHandleLease>,
38    _nosend_nosync: PhantomData<*mut ()>,
39}
40
41impl std::fmt::Debug for RngState {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        f.debug_struct("RngState").finish_non_exhaustive()
44    }
45}
46
47impl RngState {
48    /// Create a new RNG state with the given seed.
49    ///
50    /// Different seeds produce different random sequences. In multi-GPU
51    /// settings, each GPU should use a different seed.
52    pub fn new(handle: &EnteredContext<'_>, seed: u64) -> Result<Self> {
53        let mut raw: *mut cugraph_rng_state_t = ptr::null_mut();
54        let mut error: *mut cugraph_error_t = ptr::null_mut();
55
56        // SAFETY: `handle` outlives the call and `raw`/`error` are valid
57        // out-pointers for cuGraph to initialize.
58        unsafe {
59            let code = cugraph_rng_state_create(handle.handle_ptr(), seed, &mut raw, &mut error);
60            check_error(code, error, handle.device(), "RngState::new")?;
61        }
62
63        let raw = require_non_null(raw, "RngState::new", "rng_state")?;
64
65        Ok(Self {
66            raw,
67            device: handle.device(),
68            _resource_handle_lease: handle.runtime_bound_lease(),
69            _nosend_nosync: PhantomData,
70        })
71    }
72
73    /// Get the raw C pointer. For internal use by other modules.
74    pub(crate) fn as_ptr(&mut self) -> *mut cugraph_rng_state_t {
75        self.raw
76    }
77}
78
79impl Drop for RngState {
80    fn drop(&mut self) {
81        drop_device_resource(&mut self.raw, self.device, "RngState::drop", DeviceResourceKind::RngState);
82    }
83}
84
85#[cfg(test)]
86mod tests {
87
88    use super::*;
89    use crate::ResourceHandle;
90
91    #[test]
92    fn runtime_bound_rng_state_retains_native_owner_until_drop() {
93        let runtime = std::sync::Arc::new(cudf::Runtime::builder().build().expect("runtime"));
94        let handle = ResourceHandle::for_runtime(runtime).expect("runtime-bound handle");
95        let rng = {
96            let entered = handle.enter("runtime_bound_rng_state_retains_native_owner_until_drop").expect("entered");
97            RngState::new(&entered, 42).expect("rng state")
98        };
99
100        drop(handle);
101        let lease = rng._resource_handle_lease.as_ref().expect("runtime-bound rng should retain its native owner");
102        let mut entered_after_handle_drop = false;
103        lease.with_teardown_gate(|| entered_after_handle_drop = true);
104        assert!(entered_after_handle_drop, "the RNG lease must keep the native owner live after its handle drops");
105        drop(rng);
106    }
107}
cugraph/random.rs

cugraph/
random.rs
