diffusionx 0.12.0

A multi-threaded crate for random number generation and stochastic process simulation, with optional GPU acceleration.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

use crate::XResult;
use cudarc::driver::{CudaContext, CudaStream, LaunchConfig};
use std::sync::{Arc, LazyLock};

pub(crate) static CUDA_CTX: LazyLock<XResult<Arc<CudaContext>>> =
    LazyLock::new(|| Ok(CudaContext::new(0)?));

pub(crate) static CUDA_STREAM: LazyLock<XResult<Arc<CudaStream>>> =
    LazyLock::new(|| Ok(CUDA_CTX.as_ref()?.default_stream()));

pub(crate) const BM_PTX: &str = include_str!(env!("BM_KERNEL_PTX"));
pub(crate) const RANDOM_PTX: &str = include_str!(env!("RANDOM_KERNEL_PTX"));
pub(crate) const LEVY_PTX: &str = include_str!(env!("LEVY_KERNEL_PTX"));
pub(crate) const OU_PTX: &str = include_str!(env!("OU_KERNEL_PTX"));

#[inline]
pub(crate) fn config(particles: usize) -> LaunchConfig {
    let block_size = 256;
    let grid_size = particles.div_ceil(block_size);
    LaunchConfig {
        grid_dim: (grid_size as u32, 1, 1),
        block_dim: (block_size as u32, 1, 1),
        shared_mem_bytes: 0,
    }
}

/// Macro to generate GPU-accelerated moment calculation functions
macro_rules! subscribe_gpu_function {
    ($module:expr, $func_name:ident, $kernel_name:expr, ($( $param_name:ident: $param_type:ty ),+ $(,)?)) => {
        #[allow(clippy::too_many_arguments)]
        fn $func_name(
            $(
                $param_name: $param_type,
            )+
            particles: usize,
        ) -> XResult<f32> {
            let stream = $crate::gpu::CUDA_STREAM.as_ref()?;
            let kernel = $kernel_name.as_ref()?;
            let mut device_out = stream.alloc_zeros::<f32>(1)?;
            let cfg = $crate::gpu::config(particles);

            let mut rng = rand::rng();
            use rand::Rng;
            let seed: u64 = rng.random();

            let mut builder = stream.launch_builder(kernel);
            use cudarc::driver::PushKernelArg;
            builder.arg(&mut device_out);

            $(
                builder.arg(&$param_name);
            )+

            builder.arg(&particles);
            builder.arg(&seed);

            unsafe {
                builder.launch(cfg)?;
            }

            let out_host = stream.clone_dtoh(&device_out)?;
            let sum = out_host[0];
            Ok(sum / particles as f32)
        }
    };
}

/// Macro to generate GPU-accelerated central moment calculation functions
macro_rules! subscribe_central_moment_gpu_function {
    ($module:expr, $func_name:ident, $kernel_name:expr, ($( $param_name:ident: $param_type:ty ),+ $(,)?), $order_type:ty) => {
        #[allow(clippy::too_many_arguments)]
        fn $func_name(
            $(
                $param_name: $param_type,
            )+
            order: $order_type,
            particles: usize,
        ) -> XResult<f32> {
            let stream = $crate::gpu::CUDA_STREAM.as_ref()?;
            let kernel = $kernel_name.as_ref()?;
            let mut device_out = stream.alloc_zeros::<f32>(1)?;
            let cfg = $crate::gpu::config(particles);

            let mut rng = rand::rng();
            use rand::Rng;
            let seed: u64 = rng.random();

            let mean = mean($($param_name,)+ particles)?;

            let mut builder = stream.launch_builder(kernel);
            use cudarc::driver::PushKernelArg;
            builder.arg(&mut device_out);
            builder.arg(&mean);
            builder.arg(&order);

            $(
                builder.arg(&$param_name);
            )+

            builder.arg(&particles);
            builder.arg(&seed);

            unsafe {
                builder.launch(cfg)?;
            }

            let out_host = stream.clone_dtoh(&device_out)?;
            let sum = out_host[0];
            Ok(sum / particles as f32)
        }
    };
}

/// CUDA-accelerated Brownian motion estimators.
pub mod bm;
/// CUDA-accelerated Lévy process estimators.
pub mod levy;
/// CUDA-accelerated Ornstein-Uhlenbeck process estimators.
pub mod ou;
/// CUDA-accelerated random number generators.
pub mod random;