Merge branch 'feat/cuda-support' of https://github.com/rust-dd/stochastic-rs into feat/cuda-support

Author: dancixx

Cargo.toml

@@ -20,15 +20,16 @@ candle-datasets = "0.8.1"
 candle-nn = "0.8.1"
 candle-transformers = "0.8.1"
 chrono = "0.4.38"
-cudarc = { version = "0.13.9", optional = true, features = ["cuda-12080"] }
+cudarc = { version = "0.13.9", optional = true, features = ["cuda-12080", "cuda-version-from-build-system"] }
 flate2 = "1.0.34"
 gauss-quad = "0.2.1"
 impl-new-derive = "0.1.2"
 implied-vol = "1.0.0"
 indicatif = "0.17.8"
-itransformer = "1.0.1"
+# itransformer = "1.0.1"
 kendalls = "0.2.2"
 levenberg-marquardt = "0.14.0"
+libloading = { version = "0.8.6", optional = true }
 linreg = "0.2.0"
 mimalloc = { version = "0.1.43", optional = true }
 nalgebra = "0.33.2"
@@ -37,7 +38,7 @@ ndarray = { version = "0.16.1", features = [
     "matrixmultiply-threading",
     "blas",
 ] }
-ndarray-linalg = { version = "0.17.0", features = ["openblas-static"] }
+ndarray-linalg = { version = "0.17.0"}
 ndarray-npy = "0.9.1"
 ndarray-rand = "0.15.0"
 ndarray-stats = "0.6.0"
@@ -68,8 +69,8 @@ yahoo_finance_api = { version = "2.3.0", optional = true }
 [dev-dependencies]
 
 [features]
-cuda = ["dep:cudarc"]
-default = ["jemalloc"]
+cuda = ["dep:cudarc", "dep:libloading"]
+default = ["cuda"]
 jemalloc = ["dep:tikv-jemallocator"]
 malliavin = []
 mimalloc = ["dep:mimalloc"]
@@ -86,5 +87,17 @@ debug = false
 codegen-units = 1
 lto = true
 
+[target.'cfg(target_os = "macos")'.dependencies]
+ndarray-linalg = { version = "0.17.0", features = ["openblas-static"] }
+
+[target.'cfg(target_os = "macos")'.features]
+default = ["jemalloc"]
+
+[target.'cfg(target_os = "linux")'.features]
+default = ["jemalloc"]
+
+[target.'cfg(target_os = "windows")'.features]
+default = ["mimalloc"]
+
 # [package.metadata.docs.rs]
 # all-features = true

src/ai.rs

@@ -1,5 +1,5 @@
 use candle_core::Tensor;
-pub use itransformer::ITransformer;
+// pub use itransformer::ITransformer;
 
 pub mod fou;
 pub mod utils;

src/stochastic/cuda/fgn.cu

@@ -0,0 +1,106 @@
+#include <stdio.h>
+#include <cuda_runtime.h>
+#include <curand_kernel.h>
+#include <cufft.h>
+#include <cuComplex.h>
+#include <math.h>
+
+#ifdef _WIN32
+#define EXPORT __declspec(dllexport)
+#else
+#define EXPORT
+#endif
+
+__global__ void fill_random_with_eigs(
+    cuComplex *d_data,
+    const cuComplex *d_sqrt_eigs,
+    int traj_size,
+    int m,
+    unsigned long seed)
+{
+    int traj_id = blockIdx.x;
+    if (traj_id >= m)
+        return;
+
+    int idx = threadIdx.x;
+    if (idx >= traj_size)
+        return;
+
+    int data_idx = traj_id * traj_size + idx;
+
+    curandState state;
+    curand_init(seed + traj_id, idx, 0, &state);
+
+    float re = curand_normal(&state);
+    float im = curand_normal(&state);
+    cuComplex noise = make_cuComplex(re, im);
+
+    d_data[data_idx] = cuCmulf(noise, d_sqrt_eigs[idx]);
+}
+
+__global__ void scale_and_copy_to_output(
+    const cuComplex *d_data,
+    float *d_output,
+    int n,
+    int m,
+    int offset,
+    float hurst,
+    float t)
+{
+    int traj_id = blockIdx.x;
+    if (traj_id >= m)
+        return;
+
+    int idx = threadIdx.x;
+    int out_size = n - offset;
+    if (idx >= out_size)
+        return;
+
+    int data_idx = traj_id * (2 * n) + (idx + 1);
+    float scale = powf((float)n, -hurst) * powf(t, hurst);
+
+    int out_idx = traj_id * out_size + idx;
+    d_output[out_idx] = d_data[data_idx].x * scale;
+}
+
+extern "C" EXPORT void fgn_kernel(
+    const cuComplex *d_sqrt_eigs,
+    float *d_output,
+    int n,
+    int m,
+    int offset,
+    float hurst,
+    float t,
+    unsigned long seed)
+{
+    int traj_size = 2 * n;
+
+    cuComplex *d_data = nullptr;
+    cudaMalloc(&d_data, (size_t)m * traj_size * sizeof(cuComplex));
+
+    {
+        dim3 gridDim(m);
+        dim3 blockDim(traj_size);
+        fill_random_with_eigs<<<gridDim, blockDim>>>(
+            d_data, d_sqrt_eigs, traj_size, m, seed);
+        cudaDeviceSynchronize();
+    }
+
+    {
+        cufftHandle plan;
+        cufftPlan1d(&plan, traj_size, CUFFT_C2C, m);
+        cufftExecC2C(plan, d_data, d_data, CUFFT_FORWARD);
+        cudaDeviceSynchronize();
+        cufftDestroy(plan);
+    }
+
+    {
+        dim3 gridDim(m);
+        dim3 blockDim(n);
+        scale_and_copy_to_output<<<gridDim, blockDim>>>(
+            d_data, d_output, n, m, offset, hurst, t);
+        cudaDeviceSynchronize();
+    }
+
+    cudaFree(d_data);
+}

src/stochastic/cuda/fgn_linux/libfgn.exp

@@ -92,58 +92,91 @@ impl Sampling<f64> for FGN {
 
   #[cfg(feature = "cuda")]
   fn sample_cuda(&self) -> Result<Array2<f64>, Box<dyn Error>> {
-    use cudarc::driver::{CudaDevice, DeviceRepr, LaunchAsync, LaunchConfig};
-    use cudarc::nvrtc::Ptx;
+    // nvcc -shared -Xcompiler -fPIC fgn.cu -o libfgn.so -lcufft
+    // nvcc -shared fgn.cu -o fgn.dll -lcufft
+    use std::ffi::c_void;
 
-    let m = self.m.unwrap_or(1);
-    let device = CudaDevice::new(0)?;
-    device.load_ptx(Ptx::from_file("fgn.ptx"), "fgn", &["fgn_kernel"])?;
-    let f = device.get_func("fgn", "fgn_kernel").unwrap();
+    use cudarc::driver::{CudaDevice, DevicePtr, DevicePtrMut, DeviceRepr};
+
+    use libloading::{Library, Symbol};
 
     #[repr(C)]
-    #[derive(Debug, Default, Copy, Clone, PartialOrd, PartialEq)]
+    #[derive(Debug, Default, Copy, Clone)]
     pub struct cuComplex {
-      pub x: f64,
-      pub y: f64,
+      pub x: f32,
+      pub y: f32,
     }
 
     unsafe impl DeviceRepr for cuComplex {
-      fn as_kernel_param(&self) -> *mut std::ffi::c_void {
+      fn as_kernel_param(&self) -> *mut c_void {
         self as *const Self as *mut _
       }
     }
 
-    let host_sqrt_eigenvalues = self
-      .sqrt_eigenvalues
-      .map(|v| cuComplex { x: v.re, y: v.im })
-      .to_vec();
-    let d_sqrt_eigenvalues = device.htod_copy(host_sqrt_eigenvalues)?;
-    let mut d_output = device.alloc_zeros::<f64>(m * self.n)?;
+    type FgnKernelFn = unsafe extern "C" fn(
+      /* d_sqrt_eigs: */ *const cuComplex,
+      /* d_output:    */ *mut f32,
+      /* n:           */ i32,
+      /* m:           */ i32,
+      /* offset:      */ i32,
+      /* hurst:       */ f32,
+      /* t:           */ f32,
+      /* seed:        */ u64,
+    );
+
+    let lib = unsafe {
+      #[cfg(target_os = "windows")]
+      {
+        Library::new("src/stochastic/cuda/fgn_windows/fgn.dll")
+      }
+      #[cfg(target_os = "linux")]
+      {
+        Library::new("src/stochastic/cuda/fgn_linux/libfgn.so")
+      }
+    }?;
+
+    let fgn_kernel: Symbol<FgnKernelFn> = unsafe { lib.get(b"fgn_kernel") }?;
+    let device = CudaDevice::new(0)?;
+
+    let m = self.m.unwrap_or(1);
+    let n = self.n;
+    let offset = self.offset;
+    let hurst = self.hurst;
+    let t = self.t.unwrap_or(1.0);
     let seed = 42u64;
 
+    let host_sqrt_eigs: Vec<cuComplex> = self
+      .sqrt_eigenvalues
+      .iter()
+      .map(|z| cuComplex {
+        x: z.re as f32,
+        y: z.im as f32,
+      })
+      .collect();
+    let d_sqrt_eigs = device.htod_copy(host_sqrt_eigs)?;
+    let mut d_output = device.alloc_zeros::<f32>(m * (n - offset))?;
+
     unsafe {
-      f.launch(
-        LaunchConfig::for_num_elems(self.n as u32),
-        (
-          &d_sqrt_eigenvalues,
-          &mut d_output,
-          self.n as u32,
-          m as u32,
-          self.offset,
-          self.hurst,
-          self.t.unwrap_or(1.0),
-          seed,
-        ),
-      )
-    }?;
+      fgn_kernel(
+        (*d_sqrt_eigs.device_ptr()) as *const cuComplex,
+        (*d_output.device_ptr_mut()) as *mut f32,
+        n as i32,
+        m as i32,
+        offset as i32,
+        hurst as f32,
+        t as f32,
+        seed,
+      );
+    }
 
-    let host_output = device.sync_reclaim(d_output)?;
-    let mut fgn = Array2::zeros((m, self.n - self.offset));
+    let host_output: Vec<f32> = device.sync_reclaim(d_output)?;
+    let mut fgn = Array2::<f64>::zeros((m, n - offset));
     for i in 0..m {
-      for j in 0..self.n - self.offset {
-        fgn[[i, j]] = host_output[i * self.n + j];
+      for j in 0..(n - offset) {
+        fgn[[i, j]] = host_output[i * (n - offset) + j] as f64;
       }
     }
+
     Ok(fgn)
   }
 
@@ -218,6 +251,36 @@ mod tests {
     plot_1d!(fbm.sample(), "Fractional Brownian Motion (H = 0.7)");
   }
 
+  #[test]
+  #[tracing_test::traced_test]
+  #[cfg(feature = "cuda")]
+  fn fgn_cuda() {
+    let fbm = FGN::new(0.7, 500, Some(1.0), None);
+    let fgn = fbm.sample_cuda().unwrap();
+    let fgn = fgn.row(0);
+    tracing::info!("{:?}", fgn);
+    // plot_1d!(fgn, "Fractional Brownian Motion (H = 0.7)");
+    let mut path = Array1::<f64>::zeros(500);
+    for i in 1..500 {
+      path[i] += fgn[i - 1];
+    }
+    plot_1d!(path, "Fractional Brownian Motion (H = 0.7)");
+
+    let start = std::time::Instant::now();
+    for i in 0..1_000_000 {
+      let _ = fbm.sample_cuda();
+    }
+    let end = start.elapsed().as_secs();
+    tracing::info!("1000 fgn generated on cuda in: {end}");
+
+    let start = std::time::Instant::now();
+    for i in 0..1_000_000 {
+      let _ = fbm.sample();
+    }
+    let end = start.elapsed().as_secs();
+    tracing::info!("1000 fgn generated on cuda in: {end}");
+  }
+
   #[test]
   #[ignore = "Not implemented"]
   #[cfg(feature = "malliavin")]