Merge pull request #17 from rust-dd/feat/cuda-support

feat: support generating fgn on cuda

Author: dancixx

.devcontainer/devcontainer.json

@@ -0,0 +1,6 @@
+{
+  "image": "rapidsai/devcontainers:23.12-cpp-llvm16-rust-cuda12.3-ubuntu22.04",
+  "hostRequirements": { "gpu": true },
+  "workspaceFolder": "/home/coder/${localWorkspaceFolderBasename}",
+  "workspaceMount": "source=${localWorkspaceFolder},target=/home/coder/${localWorkspaceFolderBasename},type=bind"
+}

.zed/settings.json

@@ -7,7 +7,7 @@
     "rust-analyzer": {
       "initialization_options": {
         "cargo": {
-          "features": [""]
+          "features": ["cuda"]
         }
       }
     }

Cargo.toml

@@ -25,7 +25,7 @@ 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"
 linreg = "0.2.0"
@@ -67,12 +67,23 @@ yahoo_finance_api = { version = "2.3.0", optional = true }
 [dev-dependencies]
 
 [features]
-default = ["jemalloc"]
+cuda = ["dep:cudarc", "dep:libloading"]
+default = ["cuda"]
 jemalloc = ["dep:tikv-jemallocator"]
 malliavin = []
 mimalloc = ["dep:mimalloc"]
 yahoo = ["dep:time", "dep:yahoo_finance_api"]
 
+[target.'cfg(target_os = "macos")'.dependencies]
+ndarray-linalg = { version = "0.17.0", features = ["openblas-static"] }
+
+[target.'cfg(not(target_os = "macos"))'.dependencies]
+cudarc = { version = "0.13.9", optional = true, features = [
+    "cuda-12080",
+    "cuda-version-from-build-system",
+] }
+libloading = { version = "0.8.6", optional = true }
+
 [lib]
 name = "stochastic_rs"
 crate-type = ["cdylib", "lib"]

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/main.rs

@@ -1,4 +1,8 @@
+use ndarray::Array1;
 use prettytable::{format, row, Cell, Row, Table};
+use stochastic_rs::plot_1d;
+use stochastic_rs::stochastic::noise::fgn::FGN;
+use stochastic_rs::stochastic::Sampling;
 use std::error::Error;
 use std::fs::File;
 use std::io::{BufRead, BufReader};
@@ -15,36 +19,25 @@ use stochastic_rs::stats::fou_estimator::{
 // use your_crate::{FOUParameterEstimationV1, FOUParameterEstimationV2, FilterType};
 const N: usize = 10000;
 fn main() -> Result<(), Box<dyn Error>> {
-  for _ in 0..10 {
-    let mut table = Table::new();
-
-    table.add_row(Row::new(vec![
-      Cell::new("Test Case"),
-      Cell::new("Elapsed Time (ms)"),
-    ]));
-
-    let start = Instant::now();
-    let fbm = FGN::new(0.7, N, Some(1.0), None);
-    let _ = fbm.sample();
-    let duration = start.elapsed();
-    table.add_row(Row::new(vec![
-      Cell::new("Single Sample"),
-      Cell::new(&format!("{:.2?}", duration.as_millis())),
-    ]));
+  let fbm = FGN::new(0.7, 10_000, Some(1.0), Some(10000));
+  let fgn = fbm.sample_cuda().unwrap();
+  let fgn = fgn.row(0);
+  plot_1d!(fgn, "Fractional Brownian Motion (H = 0.7)");
+  let mut path = Array1::<f64>::zeros(500);
+  for i in 1..500 {
+    path[i] += path[i-1] + fgn[i];
+  }
+  plot_1d!(path, "Fractional Brownian Motion (H = 0.7)");
 
-    let start = Instant::now();
-    let fbm = FGN::new(0.7, N, Some(1.0), None);
-    for _ in 0..N {
-      let _ = fbm.sample();
-    }
-    let duration = start.elapsed();
-    table.add_row(Row::new(vec![
-      Cell::new("Repeated Samples"),
-      Cell::new(&format!("{:.2?}", duration.as_millis())),
-    ]));
+  let start = std::time::Instant::now();
+  let _ = fbm.sample_cuda();
+  let end = start.elapsed().as_millis();
+  println!("20000 fgn generated on cuda in: {end}");
 
-    table.printstd();
-  }
+  let start = std::time::Instant::now();
+    let _ = fbm.sample_par();
+  let end = start.elapsed().as_millis();
+  println!("20000 fgn generated on cuda in: {end}");
   // File paths
   // let paths = vec![
   //   "./test/kecskekut_original.txt",

src/stochastic.rs

@@ -26,6 +26,7 @@ pub mod noise;
 pub mod process;
 pub mod volatility;
 
+use std::error::Error;
 use std::sync::{Arc, Mutex};
 
 use ndarray::parallel::prelude::*;
@@ -42,6 +43,13 @@ pub trait Sampling<T: Clone + Send + Sync + Zero>: Send + Sync {
   /// Sample the process
   fn sample(&self) -> Array1<T>;
 
+  /// Sample the process with CUDA support
+  #[cfg(not(target_os = "macos"))]
+  #[cfg(feature = "cuda")]
+  fn sample_cuda(&self) -> Result<Array2<T>, Box<dyn Error>> {
+    unimplemented!()
+  }
+
   /// Parallel sampling
   fn sample_par(&self) -> Array2<T> {
     if self.m().is_none() {

src/stochastic/cuda/fgn.cu

@@ -0,0 +1,88 @@
+#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 tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= m * traj_size) return;
+    int traj_id = tid / traj_size;
+    int idx = tid % traj_size;
+    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[tid] = 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 out_size = n - offset;
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= m * out_size) return;
+    int traj_id = tid / out_size;
+    int idx = tid % out_size;
+    int data_idx = traj_id * (2 * n) + (idx + 1);
+    float scale = powf((float)n, -hurst) * powf(t, hurst);
+    d_output[tid] = 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));
+    {
+        int totalThreads = m * traj_size;
+        int blockSize = 512;
+        int gridSize = (totalThreads + blockSize - 1) / blockSize;
+        fill_random_with_eigs<<<gridSize, blockSize>>>(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);
+    }
+    {
+        int out_size = n - offset;
+        int totalThreads = m * out_size;
+        int blockSize = 512;
+        int gridSize = (totalThreads + blockSize - 1) / blockSize;
+        scale_and_copy_to_output<<<gridSize, blockSize>>>(d_data, d_output, n, m, offset, hurst, t);
+        cudaDeviceSynchronize();
+    }
+    cudaFree(d_data);
+}