diff --git a/CUDA_BENCHMARK.md b/CUDA_BENCHMARK.md new file mode 100644 index 00000000..27b58538 --- /dev/null +++ b/CUDA_BENCHMARK.md @@ -0,0 +1,91 @@ +# fGN sampling: CPU vs CUDA backends + +Head-to-head throughput of the fractional-Gaussian-noise (fGN) sampling backends, +generated by `benches/fgn_cuda_compare.rs` on an **NVIDIA RTX 4070 SUPER** (Ada, +sm_89, 12 GB). + +## Setup + +- **cpu** — `Cpu` backend: SIMD + rayon, `ndrustfft` circulant-embedding FFT. +- **cuda_native** — `CudaNative`: cudarc + cuFFT + NVRTC Philox. +- **gpu_cuda** — `CubeCl`: cubecl, hand-written shared-memory radix FFT. +- Harness: criterion, `f32`. **Throughput unit: `Melem/s` = 10⁶ generated fGN + values per second** (`Gelem/s` = 10⁹). Higher is faster. `n` = path length, + `m` = number of paths; a batch produces `n·m` values. +- Reported value = criterion's median (middle of the [lo, mid, hi] estimate). + +```bash +cargo bench --bench fgn_cuda_compare --features "cuda-native,gpu-cuda" +``` + +## Single path (`sample`, m = 1) — Melem/s + +| n | cpu | cuda_native | gpu_cuda | +|--:|--:|--:|--:| +| 256 | **129.9** | 5.6 | 2.3 | +| 512 | **129.4** | 9.1 | 3.7 | +| 1024 | **129.6** | 12.8 | 6.9 | +| 2048 | **122.4** | 25.4 | 11.5 | +| 4096 | **119.7** | 49.0 | 19.3 | +| 8192 | **116.5** | 94.5 | 27.5 | +| 16384 | 113.9 | **149.1** | 31.9 | +| 32768 | 108.2 | **312.3** | 39.7 | +| 65536 | 77.4 | **346.9** | 40.4 | +| 131072 | 74.7 | **218.8** | 23.2 | +| 262144 | 69.6 | **148.0** | 25.2 | + +`cuda_native` (cuFFT) overtakes the CPU at **n ≈ 16384**; it peaks at n = 65536 +(~**4.5× CPU**) and stays ahead to n = 262144 (~2.1× CPU). `gpu_cuda` is slowest +throughout. + +## Batch (`sample_par`, n × m) — Melem/s + +| n | m | cpu | cuda_native | gpu_cuda | +|--:|--:|--:|--:|--:| +| 1024 | 1024 | **1726** | 627 | 47 | +| 1024 | 16384 | **1751** | 557 | 49 | +| 1024 | 131072 | **1198** | 471 | 48 | +| 4096 | 16384 | 269 | **581** | 48 | +| 4096 | 65536 | 118 | **383** | 42 | +| 16384 | 4096 | **228** | 225 | 30 | +| 16384 | 16384 | 59 | **183** | 32 | +| 65536 | 1024 | **632** | 301 | 22 | +| 65536 | 4096 | **504** | 236 | 28 | +| 131072 | 512 | **382** | 288 | 22 | +| 131072 | 2048 | **369** | 221 | 25 | +| 262144 | 256 | **278** | 260 | 21 | +| 262144 | 1024 | 279 | **288** | 28 | + +## Regimes (who wins, and when) + +- **Short paths (n ≤ 1024), any m:** CPU dominates (1.2–1.75 **Gelem/s**) — the + FFTs fit cache and rayon saturates the cores; GPU launch/transfer overhead is + pure loss. +- **Many *and* long paths (n ≥ 4096 with large m):** **cuda_native wins ~3×** — + e.g. 4096×16384 (581 vs 269), 4096×65536 (383 vs 118), 16384×16384 (183 vs 59). + This is the batched-cuFFT sweet spot. +- **Very long paths, modest m (n ≥ 65536):** CPU and cuda_native are close; + the per-call host↔device copy + plan setup erodes the GPU's lead, so the winner + flips case by case (e.g. CPU wins 65536×1024; they tie at 262144×1024). +- **gpu_cuda (cubecl)** never wins — its hand-written FFT runs at a near-constant + ~20–49 Melem/s regardless of size, ~3–30× behind the others. Replacing it with + a cuFFT-class kernel is the obvious optimisation. + +## Takeaways + +- Use the **CPU** for short paths and for latency-sensitive single draws below + n≈16k. +- Use **cuda_native (cuFFT)** for large single paths (n ≥ 16k) and for + generating **many long paths at once** (the 3× regime above). +- **gpu_cuda** is not competitive today; it is kept for the cross-platform + (wgpu/Metal) path, not for raw CUDA speed. + +## cuda-oxide backend + +The experimental `cuda-oxide` backend is **not in this benchmark**: its kernels +use libdevice transcendentals (`ln`/`sqrt`/`cos`/`sin`), and `nvvmCompileProgram` +rejects the codegen's opaque-pointer NVVM IR (`parse expected type`). This +reproduces with cuda-oxide v0.1.0's **own** `math_atan` example inside its **exact +pinned nix environment** (CUDA 13.2 + LLVM 22 + nightly-2026-04-03), so it is an +upstream limitation. The backend is maintained on a separate experimental branch +until that is fixed. diff --git a/Cargo.toml b/Cargo.toml index 48d98022..fb0d0843 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,6 +10,10 @@ members = [ "stochastic-rs-viz", "stochastic-rs-py", ] +# Device-only cuda-oxide kernel crate: its own `[workspace]` root (built +# standalone with the cuda-oxide codegen backend to emit NVVM IR), excluded so +# the umbrella workspace does not claim it as a member. +exclude = ["stochastic-rs-stochastic/fgn-oxide-kernels"] [workspace.package] version = "2.3.0" @@ -61,7 +65,6 @@ owens-t = "0.1.5" # General utility anyhow = "1.0.89" -either = "1.15.0" parking_lot = "0.12.5" ordered-float = "5.0.0" chrono = "0.4.38" @@ -91,6 +94,9 @@ cubecl = { version = "0.9.0", default-features = false } cubecl-cuda = "0.9.0" cubecl-wgpu = "0.9.0" cudarc = { version = "0.19.2", features = ["cuda-12080", "cuda-version-from-build-system"] } +cuda-core = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "4a56e4220aab8ce5d085a411e7f806cebb647d14" } +cuda-device = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "4a56e4220aab8ce5d085a411e7f806cebb647d14" } +cuda-host = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "4a56e4220aab8ce5d085a411e7f806cebb647d14" } gpu-fft = "1.1.1" metal = "0.33.0" @@ -141,7 +147,9 @@ cubecl = { workspace = true, optional = true, default-features = false } cubecl-cuda = { workspace = true, optional = true } cubecl-wgpu = { workspace = true, optional = true } cudarc = { workspace = true, optional = true, features = [ "cuda-12080", "cuda-version-from-build-system", ] } -either = { workspace = true } +cuda-core = { workspace = true, optional = true } +cuda-device = { workspace = true, optional = true } +cuda-host = { workspace = true, optional = true } flate2 = { workspace = true } gpu-fft = { workspace = true, optional = true } gauss-quad = { workspace = true } @@ -223,6 +231,16 @@ name = "fgn_cuda_native" harness = false required-features = ["cuda-native"] +[[bench]] +name = "fgn_cuda_oxide" +harness = false +required-features = ["cuda-oxide-experimental"] + +[[bench]] +name = "fgn_cuda_compare" +harness = false +required-features = ["cuda-native", "gpu-cuda"] + [[bench]] name = "fgn_all_backends" harness = false @@ -318,9 +336,19 @@ name = "dual_stream_compare" harness = false required-features = ["dual-stream-rng"] +[[example]] +name = "cuda_fgn_analysis" +required-features = ["cuda-native"] + [features] ai = ["dep:stochastic-rs-ai", "stochastic-rs-ai/quant"] cuda-native = ["dep:cudarc", "cudarc/cufft", "stochastic-rs-stochastic/cuda-native"] +cuda-oxide-experimental = [ + "dep:cuda-core", + "dep:cuda-device", + "dep:cuda-host", + "stochastic-rs-stochastic/cuda-oxide-experimental", +] default = [] # Experimental: opt in to the dual-stream RNG (`SimdRngDual` / # `SimdNormalDual`). ~5–11% speedup on Ziggurat-based Normal/Exp bulk diff --git a/README.md b/README.md index b38ed507..79a91cf4 100644 --- a/README.md +++ b/README.md @@ -149,7 +149,7 @@ cargo bench --features cuda-native --bench fgn_cuda_native Single path: -| n | CPU `sample` | CUDA `sample_cuda_native(1)` | Speedup | +| n | CPU `sample` | CUDA `.on(Device::CudaNative).sample()` | Speedup | |-------:|-------------:|------------------------------:|-----------:| | 1,024 | 8.1 µs | 46 µs| 0.18× | | 4,096 | 35 µs | 84 µs| 0.42× | @@ -158,7 +158,7 @@ Single path: Batch: -| n, m | CPU `sample_par` | CUDA `sample_cuda_native` | Speedup | +| n, m | CPU `sample_par` | CUDA `.on(Device::CudaNative).sample_par` | Speedup | |--------------|------------------:|---------------------------:|---------:| | 4,096, 32 | 147 µs | 117 µs | **1.3×** | | 4,096, 512 | 1.78 ms | 2.37 ms | 0.75× | diff --git a/benches/fgn_accelerate.rs b/benches/fgn_accelerate.rs index 8e6fc519..412c694a 100644 --- a/benches/fgn_accelerate.rs +++ b/benches/fgn_accelerate.rs @@ -6,6 +6,7 @@ use criterion::Criterion; use criterion::criterion_group; use criterion::criterion_main; use stochastic_rs::simd_rng::Unseeded; +use stochastic_rs::stochastic::device::Accelerate; use stochastic_rs::stochastic::noise::fgn::Fgn; use stochastic_rs::traits::ProcessExt; @@ -17,13 +18,14 @@ fn bench_single(c: &mut Criterion) { for &n in &[1024usize, 4096, 16384, 65536] { let fgn = Fgn::new(0.7f32, n, None, Unseeded); - let _ = fgn.sample_accelerate(1); + let dev = Fgn::new(0.7f32, n, None, Unseeded).on::(); + let _ = dev.sample(); g.bench_with_input(BenchmarkId::new("cpu", n), &n, |b, _| { b.iter(|| black_box(fgn.sample())); }); g.bench_with_input(BenchmarkId::new("accelerate", n), &n, |b, _| { - b.iter(|| black_box(fgn.sample_accelerate(1).unwrap())); + b.iter(|| black_box(dev.sample())); }); } g.finish(); @@ -44,7 +46,8 @@ fn bench_batch(c: &mut Criterion) { ]; for &(n, m) in &cases { let fgn = Fgn::new(0.7f32, n, None, Unseeded); - let _ = fgn.sample_accelerate(m); + let dev = Fgn::new(0.7f32, n, None, Unseeded).on::(); + let _ = dev.sample_par(m); let label = format!("n={n},m={m}"); g.bench_with_input(BenchmarkId::new("cpu", &label), &(n, m), |b, &(_, m)| { @@ -54,7 +57,7 @@ fn bench_batch(c: &mut Criterion) { BenchmarkId::new("accelerate", &label), &(n, m), |b, &(_, m)| { - b.iter(|| black_box(fgn.sample_accelerate(m).unwrap())); + b.iter(|| black_box(dev.sample_par(m))); }, ); } diff --git a/benches/fgn_all_backends.rs b/benches/fgn_all_backends.rs index 4f6f3ea3..98a8d45a 100644 --- a/benches/fgn_all_backends.rs +++ b/benches/fgn_all_backends.rs @@ -6,6 +6,9 @@ use criterion::Criterion; use criterion::criterion_group; use criterion::criterion_main; use stochastic_rs::simd_rng::Unseeded; +use stochastic_rs::stochastic::device::Accelerate; +use stochastic_rs::stochastic::device::CubeCl; +use stochastic_rs::stochastic::device::MetalNative; use stochastic_rs::stochastic::noise::fgn::Fgn; use stochastic_rs::traits::ProcessExt; @@ -17,23 +20,26 @@ fn bench_single(c: &mut Criterion) { for &n in &[1024usize, 4096, 16384, 65536] { let fgn = Fgn::new(0.7f32, n, None, Unseeded); + let dev_gpu = Fgn::new(0.7f32, n, None, Unseeded).on::(); + let dev_metal = Fgn::new(0.7f32, n, None, Unseeded).on::(); + let dev_accel = Fgn::new(0.7f32, n, None, Unseeded).on::(); // warmup - let _ = fgn.sample_gpu(1); - let _ = fgn.sample_metal(1); - let _ = fgn.sample_accelerate(1); + let _ = dev_gpu.sample(); + let _ = dev_metal.sample(); + let _ = dev_accel.sample(); g.bench_with_input(BenchmarkId::new("cpu", n), &n, |b, _| { b.iter(|| black_box(fgn.sample())); }); g.bench_with_input(BenchmarkId::new("gpu_cubecl", n), &n, |b, _| { - b.iter(|| black_box(fgn.sample_gpu(1).unwrap())); + b.iter(|| black_box(dev_gpu.sample())); }); g.bench_with_input(BenchmarkId::new("metal", n), &n, |b, _| { - b.iter(|| black_box(fgn.sample_metal(1).unwrap())); + b.iter(|| black_box(dev_metal.sample())); }); g.bench_with_input(BenchmarkId::new("accelerate", n), &n, |b, _| { - b.iter(|| black_box(fgn.sample_accelerate(1).unwrap())); + b.iter(|| black_box(dev_accel.sample())); }); } g.finish(); @@ -54,22 +60,25 @@ fn bench_batch(c: &mut Criterion) { ]; for &(n, m) in &cases { let fgn = Fgn::new(0.7f32, n, None, Unseeded); - let _ = fgn.sample_gpu(m); - let _ = fgn.sample_metal(m); - let _ = fgn.sample_accelerate(m); + let dev_gpu = Fgn::new(0.7f32, n, None, Unseeded).on::(); + let dev_metal = Fgn::new(0.7f32, n, None, Unseeded).on::(); + let dev_accel = Fgn::new(0.7f32, n, None, Unseeded).on::(); + let _ = dev_gpu.sample_par(m); + let _ = dev_metal.sample_par(m); + let _ = dev_accel.sample_par(m); let label = format!("n={n},m={m}"); g.bench_with_input(BenchmarkId::new("cpu", &label), &m, |b, &m| { b.iter(|| black_box(fgn.sample_par(m))); }); g.bench_with_input(BenchmarkId::new("gpu_cubecl", &label), &m, |b, &m| { - b.iter(|| black_box(fgn.sample_gpu(m).unwrap())); + b.iter(|| black_box(dev_gpu.sample_par(m))); }); g.bench_with_input(BenchmarkId::new("metal", &label), &m, |b, &m| { - b.iter(|| black_box(fgn.sample_metal(m).unwrap())); + b.iter(|| black_box(dev_metal.sample_par(m))); }); g.bench_with_input(BenchmarkId::new("accelerate", &label), &m, |b, &m| { - b.iter(|| black_box(fgn.sample_accelerate(m).unwrap())); + b.iter(|| black_box(dev_accel.sample_par(m))); }); } g.finish(); diff --git a/benches/fgn_cuda_compare.rs b/benches/fgn_cuda_compare.rs new file mode 100644 index 00000000..b352df42 --- /dev/null +++ b/benches/fgn_cuda_compare.rs @@ -0,0 +1,115 @@ +//! Head-to-head fGN sampling benchmark: CPU (SIMD) vs cuda-native (cudarc + +//! cuFFT) vs gpu-cuda (cubecl). The per-backend `fgn_cuda_native` and `fgn_gpu` +//! benches each compare one GPU path against the CPU in isolation; this one +//! puts all three on the same axes so the two CUDA backends can be compared +//! directly. Requires both CUDA features and an NVIDIA GPU. +//! +//! Run: cargo bench --bench fgn_cuda_compare --features "cuda-native,gpu-cuda" +use std::hint::black_box; +use std::time::Duration; + +use criterion::BenchmarkId; +use criterion::Criterion; +use criterion::Throughput; +use criterion::criterion_group; +use criterion::criterion_main; +use stochastic_rs::simd_rng::Unseeded; +use stochastic_rs::stochastic::device::CubeCl; +use stochastic_rs::stochastic::device::CudaNative; +use stochastic_rs::stochastic::noise::fgn::Fgn; +use stochastic_rs::traits::ProcessExt; + +fn bench_single_path(c: &mut Criterion) { + let mut group = c.benchmark_group("FGN_compare_single_path"); + group.measurement_time(Duration::from_secs(3)); + group.warm_up_time(Duration::from_millis(700)); + group.sample_size(40); + + let hurst = 0.7f32; + for &n in &[ + 256usize, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144, + ] { + group.throughput(Throughput::Elements(n as u64)); + + let cpu = Fgn::new(hurst, n, None, Unseeded); + let native = Fgn::new(hurst, n, None, Unseeded).on::(); + let cubecl = Fgn::new(hurst, n, None, Unseeded).on::(); + + // Warm up the GPU contexts + kernel JIT once, outside the measurement. + let _ = native.sample(); + let _ = cubecl.sample(); + + group.bench_with_input(BenchmarkId::new("cpu", n), &n, |b, _| { + b.iter(|| black_box(cpu.sample())); + }); + group.bench_with_input(BenchmarkId::new("cuda_native", n), &n, |b, _| { + b.iter(|| black_box(native.sample())); + }); + group.bench_with_input(BenchmarkId::new("gpu_cuda", n), &n, |b, _| { + b.iter(|| black_box(cubecl.sample())); + }); + } + + group.finish(); +} + +fn bench_batch(c: &mut Criterion) { + let mut group = c.benchmark_group("FGN_compare_batch"); + group.measurement_time(Duration::from_secs(3)); + group.warm_up_time(Duration::from_millis(700)); + group.sample_size(10); + + let hurst = 0.7f32; + // Grid spans small → ~256k in n and → ~128k in m. Element count per case is + // capped near 268M (≈4.3 GB f32 device alloc) to fit the 12 GB RTX 4070 SUPER. + let cases = [ + (1024usize, 1024usize), + (1024, 16384), + (1024, 131072), + (4096, 16384), + (4096, 65536), + (16384, 4096), + (16384, 16384), + (65536, 1024), + (65536, 4096), + (131072, 512), + (131072, 2048), + (262144, 256), + (262144, 1024), + ]; + + for &(n, m) in &cases { + group.throughput(Throughput::Elements((n * m) as u64)); + let label = format!("n={n},m={m}"); + + let cpu = Fgn::new(hurst, n, None, Unseeded); + let native = Fgn::new(hurst, n, None, Unseeded).on::(); + let cubecl = Fgn::new(hurst, n, None, Unseeded).on::(); + + let _ = native.sample_par(m); + let _ = cubecl.sample_par(m); + + group.bench_with_input(BenchmarkId::new("cpu", &label), &(n, m), |b, &(_n, m)| { + b.iter(|| black_box(cpu.sample_par(m))); + }); + group.bench_with_input( + BenchmarkId::new("cuda_native", &label), + &(n, m), + |b, &(_n, m)| { + b.iter(|| black_box(native.sample_par(m))); + }, + ); + group.bench_with_input( + BenchmarkId::new("gpu_cuda", &label), + &(n, m), + |b, &(_n, m)| { + b.iter(|| black_box(cubecl.sample_par(m))); + }, + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_single_path, bench_batch); +criterion_main!(benches); diff --git a/benches/fgn_cuda_native.rs b/benches/fgn_cuda_native.rs index d0519b56..668030d4 100644 --- a/benches/fgn_cuda_native.rs +++ b/benches/fgn_cuda_native.rs @@ -6,6 +6,7 @@ use criterion::Criterion; use criterion::criterion_group; use criterion::criterion_main; use stochastic_rs::simd_rng::Unseeded; +use stochastic_rs::stochastic::device::CudaNative; use stochastic_rs::stochastic::noise::fgn::Fgn; use stochastic_rs::traits::ProcessExt; @@ -19,10 +20,9 @@ fn bench_fgn_single_path_cpu_vs_cuda_native(c: &mut Criterion) { for &n in &[1024usize, 4096, 16384, 65536] { let fgn = Fgn::new(hurst, n, None, Unseeded); + let dev = Fgn::new(hurst, n, None, Unseeded).on::(); - let _ = fgn - .sample_cuda_native(1) - .expect("cuda-native single-path warmup should succeed"); + let _ = dev.sample(); group.bench_with_input(BenchmarkId::new("cpu/sample", n), &n, |b, &_n| { b.iter(|| black_box(fgn.sample())); @@ -32,13 +32,7 @@ fn bench_fgn_single_path_cpu_vs_cuda_native(c: &mut Criterion) { BenchmarkId::new("cuda_native/sample_m1", n), &n, |b, &_n| { - b.iter(|| { - black_box( - fgn - .sample_cuda_native(1) - .expect("cuda-native single-path sampling should succeed"), - ) - }); + b.iter(|| black_box(dev.sample())); }, ); } @@ -67,10 +61,9 @@ fn bench_fgn_batch_cpu_vs_cuda_native(c: &mut Criterion) { for &(n, m) in &cases { let label = format!("n={n},m={m}"); let fgn = Fgn::new(hurst, n, None, Unseeded); + let dev = Fgn::new(hurst, n, None, Unseeded).on::(); - let _ = fgn - .sample_cuda_native(m) - .expect("cuda-native batch warmup should succeed"); + let _ = dev.sample_par(m); group.bench_with_input( BenchmarkId::new("cpu/sample_par", &label), @@ -84,13 +77,7 @@ fn bench_fgn_batch_cpu_vs_cuda_native(c: &mut Criterion) { BenchmarkId::new("cuda_native/sample", &label), &(n, m), |b, &(_n, m)| { - b.iter(|| { - black_box( - fgn - .sample_cuda_native(m) - .expect("cuda-native batch sampling should succeed"), - ) - }); + b.iter(|| black_box(dev.sample_par(m))); }, ); } diff --git a/benches/fgn_cuda_oxide.rs b/benches/fgn_cuda_oxide.rs new file mode 100644 index 00000000..7870c393 --- /dev/null +++ b/benches/fgn_cuda_oxide.rs @@ -0,0 +1,54 @@ +use std::hint::black_box; +use std::time::Duration; + +use criterion::BenchmarkId; +use criterion::Criterion; +use criterion::criterion_group; +use criterion::criterion_main; +use stochastic_rs::simd_rng::Unseeded; +use stochastic_rs::stochastic::noise::fgn::Fgn; +use stochastic_rs::traits::ProcessExt; + +fn bench_fgn_cuda_oxide(c: &mut Criterion) { + let mut group = c.benchmark_group("FGN_cuda_oxide_experimental"); + group.measurement_time(Duration::from_secs(3)); + group.warm_up_time(Duration::from_millis(700)); + group.sample_size(30); + + let cases = [ + (4096usize, 32usize), + (16384usize, 128usize), + (65536usize, 128usize), + ]; + + for &(n, m) in &cases { + let label = format!("n={n},m={m}"); + let fgn = Fgn::new(0.7f32, n, None, Unseeded); + let _ = fgn + .sample_cuda_oxide_with_module(m, "fgn_cuda_oxide") + .expect("cuda-oxide warmup should succeed"); + + group.bench_with_input(BenchmarkId::new("cpu/sample_par", &label), &m, |b, &m| { + b.iter(|| black_box(fgn.sample_par(m))); + }); + + group.bench_with_input( + BenchmarkId::new("cuda_oxide/sample", &label), + &m, + |b, &m| { + b.iter(|| { + black_box( + fgn + .sample_cuda_oxide_with_module(m, "fgn_cuda_oxide") + .expect("cuda-oxide sampling should succeed"), + ) + }); + }, + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_fgn_cuda_oxide); +criterion_main!(benches); diff --git a/benches/fgn_gpu.rs b/benches/fgn_gpu.rs index 02fe83a8..4cfde155 100644 --- a/benches/fgn_gpu.rs +++ b/benches/fgn_gpu.rs @@ -6,6 +6,7 @@ use criterion::Criterion; use criterion::criterion_group; use criterion::criterion_main; use stochastic_rs::simd_rng::Unseeded; +use stochastic_rs::stochastic::device::CubeCl; use stochastic_rs::stochastic::noise::fgn::Fgn; use stochastic_rs::traits::ProcessExt; @@ -19,23 +20,16 @@ fn bench_fgn_single_path_cpu_vs_gpu(c: &mut Criterion) { for &n in &[1024usize, 4096, 16384, 65536] { let fgn = Fgn::new(hurst, n, None, Unseeded); + let dev = Fgn::new(hurst, n, None, Unseeded).on::(); - let _ = fgn - .sample_gpu(1) - .expect("GPU single-path warmup should succeed"); + let _ = dev.sample(); group.bench_with_input(BenchmarkId::new("cpu/sample", n), &n, |b, &_n| { b.iter(|| black_box(fgn.sample())); }); group.bench_with_input(BenchmarkId::new("gpu/sample_gpu_m1", n), &n, |b, &_n| { - b.iter(|| { - black_box( - fgn - .sample_gpu(1) - .expect("GPU single-path sampling should succeed"), - ) - }); + b.iter(|| black_box(dev.sample())); }); } @@ -60,8 +54,9 @@ fn bench_fgn_batch_cpu_vs_gpu(c: &mut Criterion) { for &(n, m) in &cases { let label = format!("n={n},m={m}"); let fgn = Fgn::new(hurst, n, None, Unseeded); + let dev = Fgn::new(hurst, n, None, Unseeded).on::(); - let _ = fgn.sample_gpu(m).expect("GPU batch warmup should succeed"); + let _ = dev.sample_par(m); group.bench_with_input( BenchmarkId::new("cpu/sample_par", &label), @@ -75,13 +70,7 @@ fn bench_fgn_batch_cpu_vs_gpu(c: &mut Criterion) { BenchmarkId::new("gpu/sample_gpu", &label), &(n, m), |b, &(_n, m)| { - b.iter(|| { - black_box( - fgn - .sample_gpu(m) - .expect("GPU batch sampling should succeed"), - ) - }); + b.iter(|| black_box(dev.sample_par(m))); }, ); } diff --git a/benches/fgn_metal.rs b/benches/fgn_metal.rs index ed028787..ce160748 100644 --- a/benches/fgn_metal.rs +++ b/benches/fgn_metal.rs @@ -6,6 +6,7 @@ use criterion::Criterion; use criterion::criterion_group; use criterion::criterion_main; use stochastic_rs::simd_rng::Unseeded; +use stochastic_rs::stochastic::device::MetalNative; use stochastic_rs::stochastic::noise::fgn::Fgn; use stochastic_rs::traits::ProcessExt; @@ -17,13 +18,14 @@ fn bench_single(c: &mut Criterion) { for &n in &[1024usize, 4096, 16384, 65536] { let fgn = Fgn::new(0.7f32, n, None, Unseeded); - let _ = fgn.sample_metal(1); + let dev = Fgn::new(0.7f32, n, None, Unseeded).on::(); + let _ = dev.sample(); g.bench_with_input(BenchmarkId::new("cpu", n), &n, |b, _| { b.iter(|| black_box(fgn.sample())); }); g.bench_with_input(BenchmarkId::new("metal", n), &n, |b, _| { - b.iter(|| black_box(fgn.sample_metal(1).unwrap())); + b.iter(|| black_box(dev.sample())); }); } g.finish(); @@ -44,14 +46,15 @@ fn bench_batch(c: &mut Criterion) { ]; for &(n, m) in &cases { let fgn = Fgn::new(0.7f32, n, None, Unseeded); - let _ = fgn.sample_metal(m); + let dev = Fgn::new(0.7f32, n, None, Unseeded).on::(); + let _ = dev.sample_par(m); let label = format!("n={n},m={m}"); g.bench_with_input(BenchmarkId::new("cpu", &label), &(n, m), |b, &(_, m)| { b.iter(|| black_box(fgn.sample_par(m))); }); g.bench_with_input(BenchmarkId::new("metal", &label), &(n, m), |b, &(_, m)| { - b.iter(|| black_box(fgn.sample_metal(m).unwrap())); + b.iter(|| black_box(dev.sample_par(m))); }); } g.finish(); diff --git a/examples/cuda_fgn_analysis.rs b/examples/cuda_fgn_analysis.rs new file mode 100644 index 00000000..e75a210b --- /dev/null +++ b/examples/cuda_fgn_analysis.rs @@ -0,0 +1,160 @@ +//! CUDA FGN analysis: covariance vector comparison (CPU vs CUDA) and FBM path plots. +//! +//! Run: cargo run --example cuda_fgn_analysis --features cuda-native --release +use plotly::common::Mode; +use plotly::layout::Axis; +use plotly::{Layout, Plot, Scatter}; +use stochastic_rs::simd_rng::Unseeded; +use stochastic_rs::stochastic::device::CudaNative; +use stochastic_rs::stochastic::noise::fgn::Fgn; +use stochastic_rs::traits::ProcessExt; + +fn lag_covariance(paths: &[Vec], mean: f64, lag: usize) -> f64 { + let mut s = 0.0; + let mut c = 0usize; + for p in paths { + for i in 0..(p.len() - lag) { + s += (p[i] - mean) * (p[i + lag] - mean); + c += 1; + } + } + s / c as f64 +} + +fn unit_lag_cov(h: f64, k: usize) -> f64 { + if k == 0 { + 1.0 + } else { + 0.5 + * (((k + 1) as f64).powf(2.0 * h) - 2.0 * (k as f64).powf(2.0 * h) + + ((k - 1) as f64).powf(2.0 * h)) + } +} + +fn main() { + let h = 0.72_f64; + let n = 4096_usize; + let t = 1.0_f64; + let m = 2048_usize; + let max_lag = 20_usize; + + let fgn = Fgn::::new(h, n, Some(t), Unseeded); + let fgn_cuda = Fgn::::new(h, n, Some(t), Unseeded).on::(); + + // ── Covariance vector comparison ────────────────────────────────────────── + println!("Generating {m} CPU paths (n={n}, H={h})..."); + let cpu_paths: Vec> = (0..m).map(|_| fgn.sample().to_vec()).collect(); + let cpu_vals: Vec = cpu_paths.iter().flatten().copied().collect(); + let cpu_mean = cpu_vals.iter().sum::() / cpu_vals.len() as f64; + + println!("Generating {m} CUDA paths..."); + let cuda_paths: Vec> = fgn_cuda + .sample_par(m) + .into_iter() + .map(|p| p.to_vec()) + .collect(); + let cuda_vals: Vec = cuda_paths.iter().flatten().copied().collect(); + let cuda_mean = cuda_vals.iter().sum::() / cuda_vals.len() as f64; + + let dt = t / n as f64; + let var_theory = dt.powf(2.0 * h); + + println!( + "\n{:>4} {:>14} {:>14} {:>14} {:>10} {:>10}", + "lag", "theory", "CPU", "CUDA", "CPU/th", "CUDA/th" + ); + println!("{}", "-".repeat(78)); + + let mut lags = Vec::new(); + let mut theory_vec = Vec::new(); + let mut cpu_cov_vec = Vec::new(); + let mut cuda_cov_vec = Vec::new(); + + for k in 0..=max_lag { + let cov_theory = var_theory * unit_lag_cov(h, k); + let cov_cpu = lag_covariance(&cpu_paths, cpu_mean, k); + let cov_cuda = lag_covariance(&cuda_paths, cuda_mean, k); + + let ratio_cpu = cov_cpu / cov_theory; + let ratio_cuda = cov_cuda / cov_theory; + + println!( + "{:>4} {:>14.8e} {:>14.8e} {:>14.8e} {:>10.4} {:>10.4}", + k, cov_theory, cov_cpu, cov_cuda, ratio_cpu, ratio_cuda + ); + + lags.push(k as f64); + theory_vec.push(cov_theory); + cpu_cov_vec.push(cov_cpu); + cuda_cov_vec.push(cov_cuda); + } + + // Plot covariance vectors + let mut cov_plot = Plot::new(); + cov_plot.add_trace( + Scatter::new(lags.clone(), theory_vec) + .name("Theory") + .mode(Mode::LinesMarkers), + ); + cov_plot.add_trace( + Scatter::new(lags.clone(), cpu_cov_vec) + .name("CPU (empirical)") + .mode(Mode::Markers), + ); + cov_plot.add_trace( + Scatter::new(lags, cuda_cov_vec) + .name("CUDA (empirical)") + .mode(Mode::Markers), + ); + cov_plot.set_layout( + Layout::new() + .title(format!( + "FGN Autocovariance: CPU vs CUDA vs Theory (H={h}, n={n}, m={m})" + )) + .x_axis(Axis::new().title("Lag k")) + .y_axis(Axis::new().title("Cov(X_i, X_{i+k})")), + ); + cov_plot.write_html("fgn_covariance_comparison.html"); + println!("\nCovariance plot saved to fgn_covariance_comparison.html"); + + // ── FBM path plots from CUDA ────────────────────────────────────────────── + let fbm_n = 4096_usize; + let fbm_paths = 8_usize; + let hursts = [0.15, 0.25, 0.50, 0.75, 0.90]; + let time: Vec = (0..=fbm_n).map(|i| i as f64 / fbm_n as f64).collect(); + + for &hh in &hursts { + let mut plot = Plot::new(); + let fgn_h = Fgn::::new(hh, fbm_n, Some(1.0), Unseeded).on::(); + let batch = fgn_h.sample_par(fbm_paths); + + for (j, inc) in batch.iter().enumerate() { + let mut fbm = vec![0.0_f64; fbm_n + 1]; + for i in 0..fbm_n { + fbm[i + 1] = fbm[i] + inc[i]; + } + plot.add_trace( + Scatter::new(time.clone(), fbm) + .name(format!("path {}", j + 1)) + .opacity(0.8), + ); + } + + let regime = if hh < 0.5 { + "rough / anti-persistent" + } else if hh == 0.5 { + "standard Brownian motion" + } else { + "smooth / persistent" + }; + plot.set_layout( + Layout::new() + .title(format!("CUDA fBM H={hh:.2} ({regime})")) + .x_axis(Axis::new().title("t")) + .y_axis(Axis::new().title("B_H(t)")), + ); + let fname = format!("cuda_fbm_H{}.html", format!("{hh:.2}").replace('.', "")); + plot.write_html(&fname); + println!("Saved {fname}"); + } +} diff --git a/src/lib.rs b/src/lib.rs index 1f30f2cd..3ddd2efa 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -51,9 +51,11 @@ pub mod prelude { pub use stochastic_rs_quant::OptionStyle; pub use stochastic_rs_quant::OptionType; + pub use crate::traits::Backend; pub use crate::traits::BivariateExt; pub use crate::traits::CalibrationResult; pub use crate::traits::Calibrator; + pub use crate::traits::Cpu; pub use crate::traits::DistributionExt; pub use crate::traits::DistributionSampler; pub use crate::traits::FloatExt; diff --git a/src/traits.rs b/src/traits.rs index 6eebc174..1d8255af 100644 --- a/src/traits.rs +++ b/src/traits.rs @@ -24,6 +24,8 @@ pub use stochastic_rs_quant::traits::PricingResult; pub use stochastic_rs_quant::traits::StandardResult; pub use stochastic_rs_quant::traits::TimeExt; pub use stochastic_rs_quant::traits::ToModel; +pub use stochastic_rs_stochastic::device::Backend; +pub use stochastic_rs_stochastic::device::Cpu; pub use stochastic_rs_stochastic::traits::ComplexPathOutput; pub use stochastic_rs_stochastic::traits::CurveOutput; pub use stochastic_rs_stochastic::traits::Malliavin2DExt; diff --git a/stochastic-rs-copulas/Cargo.toml b/stochastic-rs-copulas/Cargo.toml index a2565911..8a98fe7e 100644 --- a/stochastic-rs-copulas/Cargo.toml +++ b/stochastic-rs-copulas/Cargo.toml @@ -51,6 +51,5 @@ gauss-quad = { workspace = true } quadrature = { workspace = true } owens-t = { workspace = true } plotly = { workspace = true, features = ["plotly_ndarray"] } -either = { workspace = true } pyo3 = { workspace = true, features = ["extension-module"], optional = true } numpy = { workspace = true, optional = true } diff --git a/stochastic-rs-quant/Cargo.toml b/stochastic-rs-quant/Cargo.toml index b55310c8..c8f5d66a 100644 --- a/stochastic-rs-quant/Cargo.toml +++ b/stochastic-rs-quant/Cargo.toml @@ -77,7 +77,6 @@ implied-vol = { workspace = true } chrono = { workspace = true } parking_lot = { workspace = true } anyhow = { workspace = true } -either = { workspace = true } ordered-float = { workspace = true } plotly = { workspace = true, features = ["plotly_ndarray"] } flate2 = { workspace = true } diff --git a/stochastic-rs-stats/Cargo.toml b/stochastic-rs-stats/Cargo.toml index 91c055be..ab5b72c9 100644 --- a/stochastic-rs-stats/Cargo.toml +++ b/stochastic-rs-stats/Cargo.toml @@ -65,6 +65,5 @@ chrono = { workspace = true } parking_lot = { workspace = true } ndrustfft = { workspace = true } anyhow = { workspace = true } -either = { workspace = true } pyo3 = { workspace = true, features = ["extension-module"], optional = true } numpy = { workspace = true, optional = true } diff --git a/stochastic-rs-stochastic/Cargo.toml b/stochastic-rs-stochastic/Cargo.toml index 660e1704..5027621b 100644 --- a/stochastic-rs-stochastic/Cargo.toml +++ b/stochastic-rs-stochastic/Cargo.toml @@ -25,7 +25,13 @@ openblas-static = [ gpu = ["dep:cubecl", "dep:gpu-fft"] gpu-cuda = ["gpu", "dep:cubecl-cuda"] gpu-wgpu = ["gpu", "dep:cubecl-wgpu"] -cuda-native = ["dep:cudarc"] +cuda-native = ["dep:cudarc", "cudarc/cufft"] +cuda-oxide-experimental = [ + "dep:cuda-core", + "dep:cuda-device", + "dep:cuda-host", + "dep:fgn-oxide-kernels", +] metal = ["dep:metal"] accelerate = [] python = ["dep:pyo3", "dep:numpy", "stochastic-rs-core/python", "stochastic-rs-distributions/python"] @@ -51,7 +57,6 @@ argmin-math = { workspace = true, features = ["ndarray_latest-nolinalg", "vec"] gauss-quad = { workspace = true } quadrature = { workspace = true } nalgebra = { workspace = true } -either = { workspace = true } anyhow = { workspace = true } wide = { workspace = true } sci-rs = { workspace = true } @@ -65,6 +70,14 @@ cubecl = { workspace = true, optional = true, default-features = false } cubecl-cuda = { workspace = true, optional = true } cubecl-wgpu = { workspace = true, optional = true } cudarc = { workspace = true, optional = true, features = ["cuda-12080", "cuda-version-from-build-system"] } +cuda-core = { workspace = true, optional = true } +cuda-device = { workspace = true, optional = true } +cuda-host = { workspace = true, optional = true } +# Device-only kernel crate (its own `[workspace]`); provides the `#[kernel]` +# markers the host `cuda_launch!` calls resolve against. Built as host stubs by +# a normal `cargo build`; its NVVM IR is pre-generated and embedded — see +# `noise/fgn/cuda_oxide.rs`. +fgn-oxide-kernels = { path = "fgn-oxide-kernels", optional = true } gpu-fft = { workspace = true, optional = true } metal = { workspace = true, optional = true } pyo3 = { workspace = true, features = ["extension-module"], optional = true } diff --git a/stochastic-rs-stochastic/fgn-oxide-kernels/Cargo.toml b/stochastic-rs-stochastic/fgn-oxide-kernels/Cargo.toml new file mode 100644 index 00000000..533e6acb --- /dev/null +++ b/stochastic-rs-stochastic/fgn-oxide-kernels/Cargo.toml @@ -0,0 +1,34 @@ +# Device-only kernel crate for the experimental cuda-oxide fGN backend. +# +# This crate contains ONLY `#[kernel]` device code. It is compiled two ways: +# 1. Normally (plain `cargo build`) it expands to host-side kernel markers +# that `stochastic-rs-stochastic`'s `cuda_launch!` calls resolve against. +# 2. Once, by the maintainer, with the cuda-oxide rustc codegen backend +# (`RUSTFLAGS="-Z codegen-backend=…"` + `CUDA_OXIDE_PTX_DIR`), it emits +# `fgn_oxide_kernels.ptx`. That PTX is committed and `include_str!`-ed into +# the lib so downstream users run with a plain `cargo build` — no +# `cargo oxide` / precompile step on their side. +# +# Standalone `[workspace]` so the cuda-oxide device build does not pull in the +# whole stochastic-rs workspace (host-only code can't be lowered to PTX). +[package] +name = "fgn-oxide-kernels" +version = "0.1.0" +edition = "2024" +license = "MIT" +description = "Device-side fGN kernels for the experimental cuda-oxide backend." + +# Not a workspace root: the umbrella `Cargo.toml` lists this path under +# `[workspace].exclude`, so it builds standalone for cuda-oxide PTX generation +# yet is usable as a path dependency without a "multiple workspace roots" error. + +[lib] +doctest = false + +[dependencies] +# The `#[kernel]` macro expands to code referencing all three cuda-oxide crates +# (device intrinsics + host-side registration markers), so all are required even +# though this crate is device-only. +cuda-device = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "4a56e4220aab8ce5d085a411e7f806cebb647d14" } +cuda-host = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "4a56e4220aab8ce5d085a411e7f806cebb647d14" } +cuda-core = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "4a56e4220aab8ce5d085a411e7f806cebb647d14" } diff --git a/stochastic-rs-stochastic/fgn-oxide-kernels/rust-toolchain.toml b/stochastic-rs-stochastic/fgn-oxide-kernels/rust-toolchain.toml new file mode 100644 index 00000000..13204e7a --- /dev/null +++ b/stochastic-rs-stochastic/fgn-oxide-kernels/rust-toolchain.toml @@ -0,0 +1,3 @@ +[toolchain] +channel = "nightly-2026-04-03" +components = ["rust-src", "rustc-dev", "llvm-tools"] diff --git a/stochastic-rs-stochastic/fgn-oxide-kernels/src/lib.rs b/stochastic-rs-stochastic/fgn-oxide-kernels/src/lib.rs new file mode 100644 index 00000000..28e73794 --- /dev/null +++ b/stochastic-rs-stochastic/fgn-oxide-kernels/src/lib.rs @@ -0,0 +1,235 @@ +//! Device-side fGN kernels for the experimental cuda-oxide backend. +//! +//! These are the same kernels the host launches via `cuda_launch!` in +//! `stochastic-rs-stochastic`'s `noise::fgn::cuda_oxide`. They live in their +//! own device-only crate so the cuda-oxide rustc codegen backend can lower +//! them to PTX without trying to compile the host-only code in the umbrella +//! lib. See this crate's `Cargo.toml` for the two-mode build contract. + +#![allow(clippy::missing_safety_doc)] + +use cuda_device::{DisjointSlice, kernel, thread}; + +const TAU_F32: f32 = 6.283_185_5; +const TAU_F64: f64 = 6.283_185_307_179_586; + +#[inline] +fn reverse_bits_n(mut x: usize, bits: usize) -> usize { + let mut y = 0usize; + let mut i = 0usize; + while i < bits { + y = (y << 1) | (x & 1); + x >>= 1; + i += 1; + } + y +} + +#[inline] +fn philox2x32_10(tid: usize, seed: u64, seq: u64) -> (u32, u32) { + let ctr = tid as u64 + seq; + let mut lo = ctr as u32; + let mut hi = (ctr >> 32) as u32; + let mut k = seed as u32; + let mut i = 0; + while i < 10 { + let p = 0xD251_1F53_u64 * lo as u64; + lo = ((p >> 32) as u32) ^ hi ^ k; + hi = p as u32; + k = k.wrapping_add(0x9E37_79B9); + i += 1; + } + (lo, hi) +} + +#[kernel] +pub fn gen_scale_f32( + mut data: DisjointSlice, + sqrt_eigs: &[f32], + traj_size: usize, + total_complex: usize, + seed: u64, + seq: u64, +) { + let tid = thread::index_1d().get(); + if tid < total_complex { + let (lo, hi) = philox2x32_10(tid, seed, seq); + let u1 = (lo as f32 + 0.5) * 2.328_306_4e-10; + let u2 = (hi as f32 + 0.5) * 2.328_306_4e-10; + let r = (-2.0 * u1.ln()).sqrt(); + let angle = TAU_F32 * u2; + let eig = sqrt_eigs[tid % traj_size]; + let base = 2 * tid; + unsafe { + let ptr = data.as_mut_ptr(); + *ptr.add(base) = r * angle.cos() * eig; + *ptr.add(base + 1) = r * angle.sin() * eig; + } + } +} + +#[kernel] +pub fn bit_reverse_f32(mut data: DisjointSlice, traj_size: usize, log_n: usize) { + let tid = thread::index_1d().get(); + let batch = tid / traj_size; + let local = tid % traj_size; + let rev = reverse_bits_n(local, log_n); + if local < rev { + let i = 2 * (batch * traj_size + local); + let j = 2 * (batch * traj_size + rev); + unsafe { + let ptr = data.as_mut_ptr(); + let ar = *ptr.add(i); + let ai = *ptr.add(i + 1); + *ptr.add(i) = *ptr.add(j); + *ptr.add(i + 1) = *ptr.add(j + 1); + *ptr.add(j) = ar; + *ptr.add(j + 1) = ai; + } + } +} + +#[kernel] +pub fn fft_stage_f32(mut data: DisjointSlice, traj_size: usize, half_stride: usize) { + let tid = thread::index_1d().get(); + let butterflies_per_batch = traj_size / 2; + let batch = tid / butterflies_per_batch; + let local = tid % butterflies_per_batch; + let stride = half_stride * 2; + let group = local / half_stride; + let pos = local % half_stride; + let i_complex = batch * traj_size + group * stride + pos; + let j_complex = i_complex + half_stride; + let angle = -TAU_F32 * (pos as f32) / (stride as f32); + let wr = angle.cos(); + let wi = angle.sin(); + + unsafe { + let ptr = data.as_mut_ptr(); + let i = 2 * i_complex; + let j = 2 * j_complex; + let ar = *ptr.add(i); + let ai = *ptr.add(i + 1); + let br = *ptr.add(j); + let bi = *ptr.add(j + 1); + let tr = br * wr - bi * wi; + let ti = br * wi + bi * wr; + *ptr.add(i) = ar + tr; + *ptr.add(i + 1) = ai + ti; + *ptr.add(j) = ar - tr; + *ptr.add(j + 1) = ai - ti; + } +} + +#[kernel] +pub fn extract_real_f32( + data: &[f32], + mut output: DisjointSlice, + out_size: usize, + traj_size: usize, + scale: f32, +) { + let tid = thread::index_1d(); + if let Some(out) = output.get_mut(tid) { + let flat = tid.get(); + let traj_id = flat / out_size; + let idx = flat % out_size; + *out = data[2 * (traj_id * traj_size + idx + 1)] * scale; + } +} + +#[kernel] +pub fn gen_scale_f64( + mut data: DisjointSlice, + sqrt_eigs: &[f64], + traj_size: usize, + total_complex: usize, + seed: u64, + seq: u64, +) { + let tid = thread::index_1d().get(); + if tid < total_complex { + let (lo, hi) = philox2x32_10(tid, seed, seq); + let u1 = (lo as f64 + 0.5) * 2.328_306_436_538_696_3e-10; + let u2 = (hi as f64 + 0.5) * 2.328_306_436_538_696_3e-10; + let r = (-2.0 * u1.ln()).sqrt(); + let angle = TAU_F64 * u2; + let eig = sqrt_eigs[tid % traj_size]; + let base = 2 * tid; + unsafe { + let ptr = data.as_mut_ptr(); + *ptr.add(base) = r * angle.cos() * eig; + *ptr.add(base + 1) = r * angle.sin() * eig; + } + } +} + +#[kernel] +pub fn bit_reverse_f64(mut data: DisjointSlice, traj_size: usize, log_n: usize) { + let tid = thread::index_1d().get(); + let batch = tid / traj_size; + let local = tid % traj_size; + let rev = reverse_bits_n(local, log_n); + if local < rev { + let i = 2 * (batch * traj_size + local); + let j = 2 * (batch * traj_size + rev); + unsafe { + let ptr = data.as_mut_ptr(); + let ar = *ptr.add(i); + let ai = *ptr.add(i + 1); + *ptr.add(i) = *ptr.add(j); + *ptr.add(i + 1) = *ptr.add(j + 1); + *ptr.add(j) = ar; + *ptr.add(j + 1) = ai; + } + } +} + +#[kernel] +pub fn fft_stage_f64(mut data: DisjointSlice, traj_size: usize, half_stride: usize) { + let tid = thread::index_1d().get(); + let butterflies_per_batch = traj_size / 2; + let batch = tid / butterflies_per_batch; + let local = tid % butterflies_per_batch; + let stride = half_stride * 2; + let group = local / half_stride; + let pos = local % half_stride; + let i_complex = batch * traj_size + group * stride + pos; + let j_complex = i_complex + half_stride; + let angle = -TAU_F64 * (pos as f64) / (stride as f64); + let wr = angle.cos(); + let wi = angle.sin(); + + unsafe { + let ptr = data.as_mut_ptr(); + let i = 2 * i_complex; + let j = 2 * j_complex; + let ar = *ptr.add(i); + let ai = *ptr.add(i + 1); + let br = *ptr.add(j); + let bi = *ptr.add(j + 1); + let tr = br * wr - bi * wi; + let ti = br * wi + bi * wr; + *ptr.add(i) = ar + tr; + *ptr.add(i + 1) = ai + ti; + *ptr.add(j) = ar - tr; + *ptr.add(j + 1) = ai - ti; + } +} + +#[kernel] +pub fn extract_real_f64( + data: &[f64], + mut output: DisjointSlice, + out_size: usize, + traj_size: usize, + scale: f64, +) { + let tid = thread::index_1d(); + if let Some(out) = output.get_mut(tid) { + let flat = tid.get(); + let traj_id = flat / out_size; + let idx = flat % out_size; + *out = data[2 * (traj_id * traj_size + idx + 1)] * scale; + } +} diff --git a/stochastic-rs-stochastic/src/device.rs b/stochastic-rs-stochastic/src/device.rs new file mode 100644 index 00000000..07cfdc65 --- /dev/null +++ b/stochastic-rs-stochastic/src/device.rs @@ -0,0 +1,132 @@ +//! Compile-time sampling backends. +//! +//! A process is parameterised by a backend marker `B` ([`Cpu`] is the default); +//! the [`Backend`] trait monomorphises `sample` / `sample_par` to that backend +//! with **no runtime branch**. Switch backend with the turbofish +//! `process.on::()` — the marker must be in scope, and the GPU +//! markers only exist when their feature is compiled, so selecting an +//! unavailable backend is a compile error rather than a runtime fallback. + +use ndarray::Array1; +use ndarray::parallel::prelude::*; +use stochastic_rs_core::simd_rng::SeedExt; + +use crate::noise::fgn::Fgn; +use crate::traits::FloatExt; + +/// CPU backend — the default `B` for every process. +pub struct Cpu; + +/// cudarc + cuFFT + NVRTC Philox. +#[cfg(feature = "cuda-native")] +pub struct CudaNative; + +/// cuda-oxide Rust → PTX. +#[cfg(feature = "cuda-oxide-experimental")] +pub struct CudaOxide; + +/// cubecl Rust kernels (CUDA or wgpu, per the compiled `cubecl-*` runtime). +#[cfg(feature = "gpu")] +pub struct CubeCl; + +/// Hand-written MSL via the `metal` crate. f32 only — Apple GPUs lack f64. +#[cfg(feature = "metal")] +pub struct MetalNative; + +/// Apple vDSP / AMX (FFI system framework, macOS). +#[cfg(feature = "accelerate")] +pub struct Accelerate; + +/// A compile-time fGN sampling backend. Implemented by the marker types in this +/// module; `Fgn` dispatches to `B` with zero runtime branching. +/// +/// The `Send + Sync` supertraits let a backend-parameterised process satisfy +/// the `ProcessExt: Send + Sync` bound and be shared across rayon worker +/// threads — every marker is a zero-sized unit struct, so this is free. +pub trait Backend: Sized + Send + Sync { + /// One fGN increment vector. The host-side `seed` drives the CPU path only; + /// GPU backends use the fGN's internal RNG. + fn generate(fgn: &Fgn, seed: &S2) -> Array1; + + /// `m` fGN paths in one batched call, one [`Array1`] per path. + fn generate_batch(fgn: &Fgn, m: usize) -> Vec>; + + /// Two independent fGN paths in one pass. Default: a batch of two; [`Cpu`] + /// overrides with the real/imag parts of a single circulant FFT (one FFT, + /// two independent fields — Dietrich & Newsam). The host-side `seed` drives + /// the CPU path only. + fn generate_pair( + fgn: &Fgn, + seed: &S2, + ) -> (Array1, Array1) { + let _ = seed; + let mut paths = Self::generate_batch(fgn, 2); + let second = paths.pop().expect("generate_batch(2) yields two paths"); + let first = paths.pop().expect("generate_batch(2) yields two paths"); + (first, second) + } +} + +impl Backend for Cpu { + fn generate(fgn: &Fgn, seed: &S2) -> Array1 { + fgn.sample_cpu_impl(seed) + } + + fn generate_batch(fgn: &Fgn, m: usize) -> Vec> { + (0..m).into_par_iter().map(|_| fgn.sample_cpu()).collect() + } + + fn generate_pair( + fgn: &Fgn, + seed: &S2, + ) -> (Array1, Array1) { + fgn.sample_pair_cpu_impl(seed) + } +} + +/// Generates a [`Backend`] impl for a GPU marker whose `$sampler` returns an +/// `Array2` of `m` paths. Single-path `generate` takes the first row; the +/// host-side seed is unused (GPU backends carry their own RNG). Each marker and +/// its impl are gated on the backend's feature. +macro_rules! gpu_backend { + ($feat:literal, $marker:ident => $sampler:ident) => { + #[cfg(feature = $feat)] + impl Backend for $marker { + fn generate( + fgn: &Fgn, + _seed: &S2, + ) -> Array1 { + fgn.$sampler(1).unwrap().row(0).to_owned() + } + + fn generate_batch( + fgn: &Fgn, + m: usize, + ) -> Vec> { + fgn + .$sampler(m) + .unwrap() + .outer_iter() + .map(|row| row.to_owned()) + .collect() + } + } + }; +} + +gpu_backend!("cuda-native", CudaNative => sample_cuda_native_impl); +gpu_backend!("cuda-oxide-experimental", CudaOxide => sample_cuda_oxide_impl); +gpu_backend!("gpu", CubeCl => sample_gpu_impl); +gpu_backend!("metal", MetalNative => sample_metal_impl); +gpu_backend!("accelerate", Accelerate => sample_accelerate_impl); + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cpu_marker_is_a_backend() { + fn assert_backend() {} + assert_backend::(); + } +} diff --git a/stochastic-rs-stochastic/src/diffusion/cfou.rs b/stochastic-rs-stochastic/src/diffusion/cfou.rs index 7a3a4b5a..ce51b42e 100644 --- a/stochastic-rs-stochastic/src/diffusion/cfou.rs +++ b/stochastic-rs-stochastic/src/diffusion/cfou.rs @@ -20,6 +20,8 @@ use num_complex::Complex; use stochastic_rs_core::simd_rng::SeedExt; use stochastic_rs_core::simd_rng::Unseeded; +use crate::device::Backend; +use crate::device::Cpu; use crate::noise::fgn::Fgn; use crate::traits::FloatExt; use crate::traits::ProcessExt; @@ -28,7 +30,7 @@ use crate::traits::ProcessExt; /// /// Source: /// - https://arxiv.org/abs/2406.18004 -pub struct Cfou { +pub struct Cfou { /// Hurst exponent of the driving fractional Brownian motion. pub hurst: T, /// Real part of the complex mean-reversion coefficient (`lambda > 0`). @@ -47,10 +49,10 @@ pub struct Cfou { pub t: Option, /// Seed strategy (compile-time: [`Unseeded`] or [`Deterministic`]). pub seed: S, - fgn: Fgn, + fgn: Fgn, } -impl Cfou { +impl Cfou { #[must_use] pub fn new( hurst: T, @@ -82,7 +84,9 @@ impl Cfou { } } -impl ProcessExt for Cfou { +backend_switch!([T: FloatExt, S: SeedExt] Cfou { hurst, lambda, omega, a, n, x1_0, x2_0, t, seed } via fgn); + +impl ProcessExt for Cfou { type Output = Array1>; /// Samples the complex path directly as `Z_t = X_1(t) + i X_2(t)`. @@ -95,8 +99,7 @@ impl ProcessExt for Cfou { /// - https://arxiv.org/abs/2406.18004 fn sample(&self) -> Self::Output { let dt = self.fgn.dt(); - let noise_1 = self.fgn.sample_cpu_impl(&self.seed.derive()); - let noise_2 = self.fgn.sample_cpu_impl(&self.seed.derive()); + let (noise_1, noise_2) = self.fgn.noise_pair(&self.seed.derive()); let gamma = Complex::new(self.lambda, -self.omega); let dt_c = Complex::new(dt, T::zero()); let noise_scale = (self.a * T::from_f64_fast(0.5)).sqrt(); diff --git a/stochastic-rs-stochastic/src/diffusion/fcir.rs b/stochastic-rs-stochastic/src/diffusion/fcir.rs index 965173d9..2dbccec3 100644 --- a/stochastic-rs-stochastic/src/diffusion/fcir.rs +++ b/stochastic-rs-stochastic/src/diffusion/fcir.rs @@ -8,6 +8,8 @@ use ndarray::Array1; use stochastic_rs_core::simd_rng::SeedExt; use stochastic_rs_core::simd_rng::Unseeded; +use crate::device::Backend; +use crate::device::Cpu; use crate::noise::fgn::Fgn; use crate::traits::FloatExt; use crate::traits::ProcessExt; @@ -15,7 +17,7 @@ use crate::traits::ProcessExt; /// Fractional Cox-Ingersoll-Ross (Fcir) process. /// dX(t) = theta(mu - X(t))dt + sigma * sqrt(X(t))dW^H(t) /// where X(t) is the Fcir process. -pub struct Fcir { +pub struct Fcir { /// Hurst exponent controlling roughness and long-memory. pub hurst: T, /// Long-run target level / model location parameter. @@ -34,10 +36,10 @@ pub struct Fcir { pub use_sym: Option, /// Seed strategy (compile-time: [`Unseeded`] or [`Deterministic`]). pub seed: S, - fgn: Fgn, + fgn: Fgn, } -impl Fcir { +impl Fcir { #[must_use] pub fn new( hurst: T, @@ -71,12 +73,12 @@ impl Fcir { } } -impl ProcessExt for Fcir { +impl ProcessExt for Fcir { type Output = Array1; fn sample(&self) -> Self::Output { let dt = self.fgn.dt(); - let fgn = self.fgn.sample_cpu_impl(&self.seed.derive()); + let fgn = self.fgn.noise(&self.seed.derive()); let mut fcir = Array1::::zeros(self.n); fcir[0] = self.x0.unwrap_or(T::zero()); @@ -95,6 +97,8 @@ impl ProcessExt for Fcir { } } +backend_switch!([T: FloatExt, S: SeedExt] Fcir { hurst, theta, mu, sigma, n, x0, t, use_sym, seed } via fgn); + py_process_1d!(PyFcir, Fcir, sig: (hurst, theta, mu, sigma, n, x0=None, t=None, use_sym=None, seed=None, dtype=None), params: (hurst: f64, theta: f64, mu: f64, sigma: f64, n: usize, x0: Option, t: Option, use_sym: Option) diff --git a/stochastic-rs-stochastic/src/diffusion/fgbm.rs b/stochastic-rs-stochastic/src/diffusion/fgbm.rs index 7e8c2178..89579193 100644 --- a/stochastic-rs-stochastic/src/diffusion/fgbm.rs +++ b/stochastic-rs-stochastic/src/diffusion/fgbm.rs @@ -8,11 +8,13 @@ use ndarray::Array1; use stochastic_rs_core::simd_rng::SeedExt; use stochastic_rs_core::simd_rng::Unseeded; +use crate::device::Backend; +use crate::device::Cpu; use crate::noise::fgn::Fgn; use crate::traits::FloatExt; use crate::traits::ProcessExt; -pub struct Fgbm { +pub struct Fgbm { /// Hurst exponent controlling roughness and long-memory. pub hurst: T, /// Drift / long-run mean-level parameter. @@ -27,10 +29,10 @@ pub struct Fgbm { pub t: Option, /// Seed strategy (compile-time: [`Unseeded`] or [`Deterministic`]). pub seed: S, - fgn: Fgn, + fgn: Fgn, } -impl Fgbm { +impl Fgbm { #[must_use] pub fn new(hurst: T, mu: T, sigma: T, n: usize, x0: Option, t: Option, seed: S) -> Self { assert!(n >= 2, "n must be at least 2"); @@ -48,12 +50,12 @@ impl Fgbm { } } -impl ProcessExt for Fgbm { +impl ProcessExt for Fgbm { type Output = Array1; fn sample(&self) -> Self::Output { let dt = self.fgn.dt(); - let fgn = self.fgn.sample_cpu_impl(&self.seed.derive()); + let fgn = self.fgn.noise(&self.seed.derive()); let mut fgbm = Array1::::zeros(self.n); fgbm[0] = self.x0.unwrap_or(T::zero()); @@ -66,7 +68,40 @@ impl ProcessExt for Fgbm { } } +backend_switch!([T: FloatExt, S: SeedExt] Fgbm { hurst, mu, sigma, n, x0, t, seed } via fgn); + py_process_1d!(PyFgbm, Fgbm, sig: (hurst, mu, sigma, n, x0=None, t=None, seed=None, dtype=None), params: (hurst: f64, mu: f64, sigma: f64, n: usize, x0: Option, t: Option) ); + +#[cfg(test)] +mod tests { + use stochastic_rs_core::simd_rng::Deterministic; + + use super::Fgbm; + use crate::device::Cpu; + use crate::traits::ProcessExt; + + #[test] + fn fgbm_on_cpu_matches_plain_sample() { + let mk = || { + Fgbm::::new( + 0.7, + 0.1, + 0.2, + 256, + Some(1.0), + Some(1.0), + Deterministic::new(7), + ) + }; + let plain = mk().sample(); + let on_cpu = mk().on::().sample(); + + assert_eq!(plain.len(), on_cpu.len()); + for (a, b) in plain.iter().zip(on_cpu.iter()) { + assert_eq!(a, b); + } + } +} diff --git a/stochastic-rs-stochastic/src/diffusion/fjacobi.rs b/stochastic-rs-stochastic/src/diffusion/fjacobi.rs index 8427d2d6..3c3dcf44 100644 --- a/stochastic-rs-stochastic/src/diffusion/fjacobi.rs +++ b/stochastic-rs-stochastic/src/diffusion/fjacobi.rs @@ -8,11 +8,13 @@ use ndarray::Array1; use stochastic_rs_core::simd_rng::SeedExt; use stochastic_rs_core::simd_rng::Unseeded; +use crate::device::Backend; +use crate::device::Cpu; use crate::noise::fgn::Fgn; use crate::traits::FloatExt; use crate::traits::ProcessExt; -pub struct FJacobi { +pub struct FJacobi { /// Hurst exponent controlling roughness and long-memory. pub hurst: T, /// Model shape / loading parameter. @@ -29,10 +31,10 @@ pub struct FJacobi { pub t: Option, /// Seed strategy (compile-time: [`Unseeded`] or [`Deterministic`]). pub seed: S, - fgn: Fgn, + fgn: Fgn, } -impl FJacobi { +impl FJacobi { #[must_use] pub fn new( hurst: T, @@ -64,12 +66,12 @@ impl FJacobi { } } -impl ProcessExt for FJacobi { +impl ProcessExt for FJacobi { type Output = Array1; fn sample(&self) -> Self::Output { let dt = self.fgn.dt(); - let fgn = self.fgn.sample_cpu_impl(&self.seed.derive()); + let fgn = self.fgn.noise(&self.seed.derive()); let mut fjacobi = Array1::::zeros(self.n); fjacobi[0] = self.x0.unwrap_or(T::zero()); @@ -90,6 +92,8 @@ impl ProcessExt for FJacobi { } } +backend_switch!([T: FloatExt, S: SeedExt] FJacobi { hurst, alpha, beta, sigma, n, x0, t, seed } via fgn); + py_process_1d!(PyFJacobi, FJacobi, sig: (hurst, alpha, beta, sigma, n, x0=None, t=None, seed=None, dtype=None), params: (hurst: f64, alpha: f64, beta: f64, sigma: f64, n: usize, x0: Option, t: Option) diff --git a/stochastic-rs-stochastic/src/diffusion/fou.rs b/stochastic-rs-stochastic/src/diffusion/fou.rs index 8b79834c..623bb120 100644 --- a/stochastic-rs-stochastic/src/diffusion/fou.rs +++ b/stochastic-rs-stochastic/src/diffusion/fou.rs @@ -8,11 +8,13 @@ use ndarray::Array1; use stochastic_rs_core::simd_rng::SeedExt; use stochastic_rs_core::simd_rng::Unseeded; +use crate::device::Backend; +use crate::device::Cpu; use crate::noise::fgn::Fgn; use crate::traits::FloatExt; use crate::traits::ProcessExt; -pub struct Fou { +pub struct Fou { /// Hurst exponent controlling roughness and long-memory. pub hurst: T, /// Mean-reversion speed. @@ -29,10 +31,10 @@ pub struct Fou { pub t: Option, /// Seed strategy (compile-time: [`Unseeded`] or [`Deterministic`]). pub seed: S, - fgn: Fgn, + fgn: Fgn, } -impl Fou { +impl Fou { #[must_use] pub fn new( hurst: T, @@ -60,12 +62,12 @@ impl Fou { } } -impl ProcessExt for Fou { +impl ProcessExt for Fou { type Output = Array1; fn sample(&self) -> Self::Output { let dt = self.fgn.dt(); - let fgn = self.fgn.sample_cpu_impl(&self.seed.derive()); + let fgn = self.fgn.noise(&self.seed.derive()); let mut fou = Array1::::zeros(self.n); fou[0] = self.x0.unwrap_or(T::zero()); @@ -78,6 +80,8 @@ impl ProcessExt for Fou { } } +backend_switch!([T: FloatExt, S: SeedExt] Fou { hurst, theta, mu, sigma, n, x0, t, seed } via fgn); + py_process_1d!(PyFou, Fou, sig: (hurst, theta, mu, sigma, n, x0=None, t=None, seed=None, dtype=None), params: (hurst: f64, theta: f64, mu: f64, sigma: f64, n: usize, x0: Option, t: Option) diff --git a/stochastic-rs-stochastic/src/jump/jump_fou.rs b/stochastic-rs-stochastic/src/jump/jump_fou.rs index bfc616d4..ba5fb26a 100644 --- a/stochastic-rs-stochastic/src/jump/jump_fou.rs +++ b/stochastic-rs-stochastic/src/jump/jump_fou.rs @@ -11,12 +11,14 @@ use stochastic_rs_core::simd_rng::Deterministic; use stochastic_rs_core::simd_rng::SeedExt; use stochastic_rs_core::simd_rng::Unseeded; +use crate::device::Backend; +use crate::device::Cpu; use crate::noise::fgn::Fgn; use crate::process::cpoisson::CompoundPoisson; use crate::traits::FloatExt; use crate::traits::ProcessExt; -pub struct JumpFou +pub struct JumpFou where T: FloatExt, D: Distribution + Send + Sync, @@ -29,11 +31,11 @@ where pub x0: Option, pub t: Option, pub cpoisson: CompoundPoisson, - fgn: Fgn, + fgn: Fgn, pub seed: S, } -impl JumpFou +impl JumpFou where T: FloatExt, D: Distribution + Send + Sync, @@ -66,7 +68,7 @@ where } } -impl ProcessExt for JumpFou +impl ProcessExt for JumpFou where T: FloatExt, D: Distribution + Send + Sync, @@ -92,6 +94,9 @@ where } } +backend_switch!([T, D, S: SeedExt] JumpFou { hurst, theta, mu, sigma, n, x0, t, cpoisson, seed } via fgn + where T: FloatExt, D: Distribution + Send + Sync); + #[cfg(feature = "python")] #[pyo3::prelude::pyclass] pub struct PyJumpFou { diff --git a/stochastic-rs-stochastic/src/jump/jump_fou_custom.rs b/stochastic-rs-stochastic/src/jump/jump_fou_custom.rs index 398545c0..2304252d 100644 --- a/stochastic-rs-stochastic/src/jump/jump_fou_custom.rs +++ b/stochastic-rs-stochastic/src/jump/jump_fou_custom.rs @@ -9,11 +9,13 @@ use rand_distr::Distribution; use stochastic_rs_core::simd_rng::SeedExt; use stochastic_rs_core::simd_rng::Unseeded; +use crate::device::Backend; +use crate::device::Cpu; use crate::noise::fgn::Fgn; use crate::traits::FloatExt; use crate::traits::ProcessExt; -pub struct JumpFOUCustom +pub struct JumpFOUCustom where T: FloatExt, D: Distribution + Send + Sync, @@ -27,11 +29,11 @@ where pub t: Option, pub jump_times: D, pub jump_sizes: D, - fgn: Fgn, + fgn: Fgn, pub seed: S, } -impl JumpFOUCustom +impl JumpFOUCustom where T: FloatExt, D: Distribution + Send + Sync, @@ -66,7 +68,7 @@ where } } -impl ProcessExt for JumpFOUCustom +impl ProcessExt for JumpFOUCustom where T: FloatExt, D: Distribution + Send + Sync, @@ -114,6 +116,9 @@ where } } +backend_switch!([T, D, S: SeedExt] JumpFOUCustom { hurst, theta, mu, sigma, n, x0, t, jump_times, jump_sizes, seed } via fgn + where T: FloatExt, D: Distribution + Send + Sync); + #[cfg(test)] mod tests { use rand_distr::Distribution; diff --git a/stochastic-rs-stochastic/src/lib.rs b/stochastic-rs-stochastic/src/lib.rs index d1602f1e..1139b096 100644 --- a/stochastic-rs-stochastic/src/lib.rs +++ b/stochastic-rs-stochastic/src/lib.rs @@ -19,6 +19,8 @@ pub use stochastic_rs_distributions as distributions; pub use stochastic_rs_distributions::traits::DistributionExt; pub use stochastic_rs_distributions::traits::SimdFloatExt; +pub use crate::device::Backend; +pub use crate::device::Cpu; pub use crate::traits::Malliavin2DExt; pub use crate::traits::MalliavinExt; pub use crate::traits::ProcessExt; @@ -26,6 +28,7 @@ pub use crate::traits::ProcessExt; pub mod aliases; pub mod autoregressive; pub mod correlation; +pub mod device; pub mod diffusion; pub mod interest; pub mod isonormal; diff --git a/stochastic-rs-stochastic/src/macros.rs b/stochastic-rs-stochastic/src/macros.rs index 82b02792..e0f8463a 100644 --- a/stochastic-rs-stochastic/src/macros.rs +++ b/stochastic-rs-stochastic/src/macros.rs @@ -240,3 +240,44 @@ macro_rules! py_process_2d { macro_rules! py_process_2d { ($($tt:tt)*) => {}; } + +/// Compile-time backend switch: generates the whole `impl` block whose +/// `on::()` re-types a process to sample on another +/// [`crate::device::Backend`] with zero runtime cost. The marker params `B` +/// (source) and `B2` (target) are appended automatically — pass the remaining +/// generics in `[..]`, the type without `B`, the moved fields, the storage +/// form, then any trailing `where` bounds. +/// +/// Storage form: +/// - `via fgn` — backend carried through an inner `fgn: Fgn<_, _, B>` field. +/// - `via phantom` — backend carried through a `_backend: PhantomData` field. +macro_rules! backend_switch { + ( + [$($gen:tt)*] $ty:ident<$($targ:ident),+ $(,)?> { $($field:ident),* $(,)? } via fgn + $(where $($wc:tt)*)? + ) => { + impl<$($gen)*, B> $ty<$($targ),+, B> $(where $($wc)*)? { + /// Re-type this process to sample on backend `B2` (compile-time, zero runtime cost). + pub fn on(self) -> $ty<$($targ),+, B2> { + $ty { + $($field: self.$field,)* + fgn: self.fgn.on::(), + } + } + } + }; + ( + [$($gen:tt)*] $ty:ident<$($targ:ident),+ $(,)?> { $($field:ident),* $(,)? } via phantom + $(where $($wc:tt)*)? + ) => { + impl<$($gen)*, B> $ty<$($targ),+, B> $(where $($wc)*)? { + /// Re-type this process to sample on backend `B2` (compile-time, zero runtime cost). + pub fn on(self) -> $ty<$($targ),+, B2> { + $ty { + $($field: self.$field,)* + _backend: ::std::marker::PhantomData, + } + } + } + }; +} diff --git a/stochastic-rs-stochastic/src/noise/cfgns.rs b/stochastic-rs-stochastic/src/noise/cfgns.rs index 2d805289..1788a2cd 100644 --- a/stochastic-rs-stochastic/src/noise/cfgns.rs +++ b/stochastic-rs-stochastic/src/noise/cfgns.rs @@ -10,10 +10,12 @@ use stochastic_rs_core::simd_rng::SeedExt; use stochastic_rs_core::simd_rng::Unseeded; use super::fgn::Fgn; +use crate::device::Backend; +use crate::device::Cpu; use crate::traits::FloatExt; use crate::traits::ProcessExt; -pub struct Cfgns { +pub struct Cfgns { /// Hurst exponent controlling roughness and long-memory. pub hurst: T, /// Instantaneous correlation parameter. @@ -24,10 +26,10 @@ pub struct Cfgns { pub t: Option, /// Seed strategy (compile-time: [`Unseeded`] or [`Deterministic`]). pub seed: S, - fgn: Fgn, + fgn: Fgn, } -impl Cfgns { +impl Cfgns { pub fn new(hurst: T, rho: T, n: usize, t: Option, seed: S) -> Self { assert!( (T::zero()..=T::one()).contains(&hurst), @@ -49,19 +51,18 @@ impl Cfgns { } } -impl Cfgns { +impl Cfgns { /// Sample with an explicit seed, used by callers like Cfbms. pub fn sample_with_seed(&self, seed: u64) -> [Array1; 2] { self.sample_impl(&Deterministic::new(seed)) } /// Core sampling — monomorphised per seed strategy, zero runtime branching. + /// Uses one paired fGN pass (real/imag of a single circulant FFT) for the two + /// independent fields; on a GPU backend they come from a batch of two. #[inline] pub(crate) fn sample_impl(&self, seed: &S2) -> [Array1; 2] { - let child1 = seed.derive(); - let child2 = seed.derive(); - let fgn1 = self.fgn.sample_cpu_impl(&child1); - let z = self.fgn.sample_cpu_impl(&child2); + let (fgn1, z) = self.fgn.noise_pair(seed); let c = (T::one() - self.rho.powi(2)).sqrt(); let mut fgn2 = Array1::zeros(self.n); for i in 0..self.n { @@ -71,7 +72,7 @@ impl Cfgns { } } -impl ProcessExt for Cfgns { +impl ProcessExt for Cfgns { type Output = [Array1; 2]; fn sample(&self) -> Self::Output { @@ -79,6 +80,8 @@ impl ProcessExt for Cfgns { } } +backend_switch!([T: FloatExt, S: SeedExt] Cfgns { hurst, rho, n, t, seed } via fgn); + py_process_2x1d!(PyCfgns, Cfgns, sig: (hurst, rho, n, t=None, seed=None, dtype=None), params: (hurst: f64, rho: f64, n: usize, t: Option) diff --git a/stochastic-rs-stochastic/src/noise/fgn.rs b/stochastic-rs-stochastic/src/noise/fgn.rs index ec5cddf3..7074bf75 100644 --- a/stochastic-rs-stochastic/src/noise/fgn.rs +++ b/stochastic-rs-stochastic/src/noise/fgn.rs @@ -9,6 +9,8 @@ mod accelerate; mod core; #[cfg(feature = "cuda-native")] mod cuda_native; +#[cfg(feature = "cuda-oxide-experimental")] +mod cuda_oxide; #[cfg(feature = "gpu")] mod gpu; #[cfg(feature = "metal")] @@ -18,47 +20,23 @@ mod python; pub use core::Fgn; -#[cfg(any( - feature = "gpu", - feature = "cuda-native", - feature = "accelerate", - feature = "metal" -))] -use anyhow::Result; -#[cfg(any( - feature = "gpu", - feature = "cuda-native", - feature = "accelerate", - feature = "metal" -))] -use either::Either; use ndarray::Array1; -#[cfg(any( - feature = "gpu", - feature = "cuda-native", - feature = "accelerate", - feature = "metal" -))] -use ndarray::Array2; #[cfg(feature = "python")] pub use python::PyFgn; use stochastic_rs_core::simd_rng::SeedExt; +use crate::device::Backend; use crate::traits::FloatExt; use crate::traits::ProcessExt; -impl Fgn { - /// Sample two independent fGn paths in a single FFT / RNG pass. +impl Fgn { + /// Sample two independent fGn paths in a single FFT / RNG pass (CPU). /// - /// Both paths have the same covariance structure as [`sample`](Self::sample) - /// and are statistically independent — Dietrich & Newsam (1997) and - /// Kroese & Botev (2013 §2.2) identify the real and imaginary parts of - /// the circulant-embedding FFT output as two independent realisations of - /// the target Gaussian field. - /// - /// For Monte-Carlo inner loops this ~halves wall time per path vs. - /// two back-to-back [`sample`](Self::sample) calls (one FFT / RNG pass - /// reused for both outputs). + /// Both paths have the same covariance structure as `sample` and are + /// statistically independent — Dietrich & Newsam (1997) and Kroese & Botev + /// (2013 §2.2) identify the real and imaginary parts of the circulant- + /// embedding FFT output as two independent realisations of the target field. + /// For Monte-Carlo inner loops this ~halves wall time per path. pub fn sample_pair(&self) -> (Array1, Array1) { self.sample_pair_cpu() } @@ -69,30 +47,16 @@ impl Fgn { } } -impl ProcessExt for Fgn { +impl ProcessExt for Fgn { type Output = Array1; fn sample(&self) -> Self::Output { - self.sample_cpu() - } - - #[cfg(feature = "gpu")] - fn sample_gpu(&self, m: usize) -> Result, Array2>> { - self.sample_gpu_impl(m) - } - - #[cfg(feature = "cuda-native")] - fn sample_cuda_native(&self, m: usize) -> Result, Array2>> { - self.sample_cuda_native_impl(m) - } - - #[cfg(feature = "accelerate")] - fn sample_accelerate(&self, m: usize) -> Result, Array2>> { - self.sample_accelerate_impl(m) + B::generate(self, &self.seed) } - #[cfg(feature = "metal")] - fn sample_metal(&self, m: usize) -> Result, Array2>> { - self.sample_metal_impl(m) + /// The `m` paths are generated in **one batched backend call** (a single FFT + /// plan over the whole batch). + fn sample_par(&self, m: usize) -> Vec { + B::generate_batch(self, m) } } diff --git a/stochastic-rs-stochastic/src/noise/fgn/accelerate.rs b/stochastic-rs-stochastic/src/noise/fgn/accelerate.rs index ee2c7f11..9733e2b7 100644 --- a/stochastic-rs-stochastic/src/noise/fgn/accelerate.rs +++ b/stochastic-rs-stochastic/src/noise/fgn/accelerate.rs @@ -8,11 +8,7 @@ use std::any::TypeId; use std::ffi::c_void; use anyhow::Result; -use either::Either; -use ndarray::Array1; use ndarray::Array2; -use rayon::iter::IntoParallelIterator; -use rayon::iter::ParallelIterator; use stochastic_rs_core::simd_rng::SeedExt; use super::Fgn; @@ -48,7 +44,7 @@ fn sample_f32( hurst: f64, t: f64, seed: &S, -) -> Result, Array2>> { +) -> Result> { let traj_size = 2 * n; let out_size = n - offset; let scale = (out_size.max(1) as f32).powf(-(hurst as f32)) * (t as f32).powf(hurst as f32); @@ -101,10 +97,7 @@ fn sample_f32( } let fgn = arr2_f32::(&output, m, out_size); - if m == 1 { - return Ok(Either::Left(fgn.row(0).to_owned())); - } - Ok(Either::Right(fgn)) + Ok(fgn) } fn arr2_f32(data: &[f32], m: usize, cols: usize) -> Array2 { @@ -122,8 +115,8 @@ fn arr2_f32(data: &[f32], m: usize, cols: usize) -> Array2 { } } -impl Fgn { - pub(crate) fn sample_accelerate_impl(&self, m: usize) -> Result, Array2>> { +impl Fgn { + pub(crate) fn sample_accelerate_impl(&self, m: usize) -> Result> { let n = self.n; let offset = self.offset; let hurst = self.hurst.to_f64().unwrap(); diff --git a/stochastic-rs-stochastic/src/noise/fgn/core.rs b/stochastic-rs-stochastic/src/noise/fgn/core.rs index 3b996250..e1ecc255 100644 --- a/stochastic-rs-stochastic/src/noise/fgn/core.rs +++ b/stochastic-rs-stochastic/src/noise/fgn/core.rs @@ -4,6 +4,7 @@ //! \operatorname{Cov}(\Delta B_i^H,\Delta B_j^H)=\tfrac12\left(|k+1|^{2H}-2|k|^{2H}+|k-1|^{2H}\right),\ k=i-j //! $$ //! +use std::marker::PhantomData; use std::sync::Arc; use ndarray::prelude::*; @@ -14,9 +15,11 @@ use stochastic_rs_core::simd_rng::Deterministic; use stochastic_rs_core::simd_rng::SeedExt; use stochastic_rs_core::simd_rng::Unseeded; +use crate::device::Backend; +use crate::device::Cpu; use crate::traits::FloatExt; -pub struct Fgn { +pub struct Fgn { /// Hurst exponent controlling roughness and long-memory. pub hurst: T, /// Internal FFT length (power-of-two padded). @@ -33,16 +36,11 @@ pub struct Fgn { pub fft_handler: Arc>, /// Seed strategy (compile-time: [`Unseeded`] or [`Deterministic`]). pub seed: S, + /// Compile-time sampling backend marker (default [`Cpu`]). + _backend: PhantomData, } -impl Fgn { - pub fn dt(&self) -> T { - let step_count = self.out_len.max(1); - self.t.unwrap_or(T::one()) / T::from_usize_(step_count) - } -} - -impl Fgn { +impl Fgn { #[must_use] pub fn new(hurst: T, n: usize, t: Option, seed: S) -> Self { assert!( @@ -102,11 +100,17 @@ impl Fgn { sqrt_eigenvalues: Arc::new(sqrt_eigenvalues), fft_handler, seed, + _backend: PhantomData, } } } -impl Fgn { +impl Fgn { + pub fn dt(&self) -> T { + let step_count = self.out_len.max(1); + self.t.unwrap_or(T::one()) / T::from_usize_(step_count) + } + /// Sample fGn using a specific deterministic seed. pub fn sample_cpu_with_seed(&self, seed: u64) -> Array1 { self.sample_cpu_impl(&Deterministic::new(seed)) @@ -187,6 +191,26 @@ impl Fgn { } } +backend_switch!([T: FloatExt, S: SeedExt] Fgn { hurst, n, t, offset, out_len, scale, sqrt_eigenvalues, fft_handler, seed } via phantom); + +impl Fgn { + /// One fGN increment vector on backend `B`. The host-side `seed` drives the + /// CPU path only; GPU backends use the fGN's internal RNG. + pub(crate) fn noise(&self, seed: &S2) -> Array1 { + B::generate(self, seed) + } + + /// `m` fGN paths in one batched `B` call, one [`Array1`] per path. + pub(crate) fn noise_batch(&self, m: usize) -> Vec> { + B::generate_batch(self, m) + } + + /// Two independent fGN paths in one pass on backend `B`. + pub(crate) fn noise_pair(&self, seed: &S2) -> (Array1, Array1) { + B::generate_pair(self, seed) + } +} + #[cfg(test)] mod tests { use stochastic_rs_core::simd_rng::Unseeded; diff --git a/stochastic-rs-stochastic/src/noise/fgn/cuda_native/sampler.rs b/stochastic-rs-stochastic/src/noise/fgn/cuda_native/sampler.rs index 10760d3b..18b349a6 100644 --- a/stochastic-rs-stochastic/src/noise/fgn/cuda_native/sampler.rs +++ b/stochastic-rs-stochastic/src/noise/fgn/cuda_native/sampler.rs @@ -4,8 +4,6 @@ use std::sync::atomic::Ordering; use anyhow::Result; use cudarc::cufft; use cudarc::driver::*; -use either::Either; -use ndarray::Array1; use ndarray::Array2; use stochastic_rs_core::simd_rng::SeedExt; @@ -30,7 +28,7 @@ fn sample_f32( hurst: f64, t: f64, seed_src: &S, -) -> Result, Array2>> { +) -> Result> { let hurst_bits = hurst.to_bits(); let t_bits = t.to_bits(); let out_size = n - offset; @@ -135,10 +133,7 @@ fn sample_f32( drop(sized); let fgn = array2_from_vec_f32::(host, m, out_size); - if m == 1 { - return Ok(Either::Left(fgn.row(0).to_owned())); - } - Ok(Either::Right(fgn)) + Ok(fgn) } fn sample_f64( @@ -149,7 +144,7 @@ fn sample_f64( hurst: f64, t: f64, seed_src: &S, -) -> Result, Array2>> { +) -> Result> { let hurst_bits = hurst.to_bits(); let t_bits = t.to_bits(); let out_size = n - offset; @@ -254,14 +249,11 @@ fn sample_f64( drop(sized); let fgn = array2_from_vec_f64::(host, m, out_size); - if m == 1 { - return Ok(Either::Left(fgn.row(0).to_owned())); - } - Ok(Either::Right(fgn)) + Ok(fgn) } -impl Fgn { - pub(crate) fn sample_cuda_native_impl(&self, m: usize) -> Result, Array2>> { +impl Fgn { + pub(crate) fn sample_cuda_native_impl(&self, m: usize) -> Result> { let n = self.n; let offset = self.offset; let hurst = self.hurst.to_f64().unwrap(); diff --git a/stochastic-rs-stochastic/src/noise/fgn/cuda_native/tests.rs b/stochastic-rs-stochastic/src/noise/fgn/cuda_native/tests.rs index c499c8d7..e14b3064 100644 --- a/stochastic-rs-stochastic/src/noise/fgn/cuda_native/tests.rs +++ b/stochastic-rs-stochastic/src/noise/fgn/cuda_native/tests.rs @@ -1,5 +1,5 @@ use super::super::Fgn; -use crate::traits::ProcessExt; +use stochastic_rs_core::simd_rng::Unseeded; fn lag_covariance(paths: &[Vec], mean: f64, lag: usize) -> f64 { let mut s = 0.0; @@ -15,42 +15,42 @@ fn lag_covariance(paths: &[Vec], mean: f64, lag: usize) -> f64 { #[test] fn cuda_native_single_path_shape() { - let fgn = Fgn::::new(0.7, 1024, Some(1.0)); + let fgn = Fgn::::new(0.7, 1024, Some(1.0), Unseeded); let result = fgn - .sample_cuda_native(1) + .sample_cuda_native_impl(1) .expect("single path should succeed"); - let path = result.left().expect("m=1 should return Left(Array1)"); - assert_eq!(path.len(), 1024); + assert_eq!(result.shape(), &[1, 1024]); } #[test] fn cuda_native_batch_shape() { - let fgn = Fgn::::new(0.7, 512, Some(1.0)); + let fgn = Fgn::::new(0.7, 512, Some(1.0), Unseeded); let m = 64; - let result = fgn.sample_cuda_native(m).expect("batch should succeed"); - let batch = result.right().expect("m>1 should return Right(Array2)"); + let batch = fgn + .sample_cuda_native_impl(m) + .expect("batch should succeed"); assert_eq!(batch.shape(), &[m, 512]); } #[test] fn cuda_native_f32_works() { - let fgn = Fgn::::new(0.7, 1024, Some(1.0)); - let result = fgn.sample_cuda_native(4).expect("f32 should succeed"); - let batch = result.right().expect("m>1 should return Right(Array2)"); + let fgn = Fgn::::new(0.7, 1024, Some(1.0), Unseeded); + let batch = fgn.sample_cuda_native_impl(4).expect("f32 should succeed"); assert_eq!(batch.shape(), &[4, 1024]); } #[test] fn cuda_native_non_power_of_two_n() { - let fgn = Fgn::::new(0.7, 3000, Some(1.0)); - let result = fgn.sample_cuda_native(8).expect("non-pot n should work"); - let batch = result.right().expect("batch"); + let fgn = Fgn::::new(0.7, 3000, Some(1.0), Unseeded); + let batch = fgn + .sample_cuda_native_impl(8) + .expect("non-pot n should work"); assert_eq!(batch.shape(), &[8, 3000]); } #[test] fn cuda_native_eigenvalues_structural() { - let fgn = Fgn::::new(0.72, 2048, Some(1.0)); + let fgn = Fgn::::new(0.72, 2048, Some(1.0), Unseeded); let eigs = &*fgn.sqrt_eigenvalues; assert_eq!(eigs.len(), 2 * fgn.n); @@ -77,7 +77,7 @@ fn cuda_native_eigenvalues_structural() { #[test] fn cuda_native_scale_matches_cpu() { for &n in &[1024_usize, 3000, 4096] { - let fgn = Fgn::::new(0.7, n, Some(2.0)); + let fgn = Fgn::::new(0.7, n, Some(2.0), Unseeded); let cpu_scale = fgn.scale; let out_size = fgn.n - fgn.offset; @@ -97,7 +97,7 @@ fn cuda_native_variance_matches_cpu() { let n = 2048_usize; let t = 1.0_f64; let m = 1024_usize; - let fgn = Fgn::::new(h, n, Some(t)); + let fgn = Fgn::::new(h, n, Some(t), Unseeded); let cpu_paths: Vec> = (0..m).map(|_| fgn.sample_cpu().to_vec()).collect(); let cpu_vals: Vec = cpu_paths.iter().flatten().copied().collect(); @@ -105,10 +105,9 @@ fn cuda_native_variance_matches_cpu() { let cpu_var = cpu_vals.iter().map(|x| (x - cpu_mean).powi(2)).sum::() / cpu_vals.len() as f64; - let cuda_result = fgn - .sample_cuda_native(m) + let cuda_batch = fgn + .sample_cuda_native_impl(m) .expect("cuda batch should succeed"); - let cuda_batch = cuda_result.right().expect("batch"); let cuda_vals: Vec = cuda_batch.iter().copied().collect(); let cuda_mean = cuda_vals.iter().sum::() / cuda_vals.len() as f64; let cuda_var = cuda_vals @@ -130,7 +129,7 @@ fn cuda_native_covariance_structure_matches_cpu() { let n = 2048_usize; let t = 1.0_f64; let m = 1024_usize; - let fgn = Fgn::::new(h, n, Some(t)); + let fgn = Fgn::::new(h, n, Some(t), Unseeded); let cpu_paths: Vec> = (0..m).map(|_| fgn.sample_cpu().to_vec()).collect(); let cpu_vals: Vec = cpu_paths.iter().flatten().copied().collect(); @@ -138,10 +137,9 @@ fn cuda_native_covariance_structure_matches_cpu() { let cpu_cov1 = lag_covariance(&cpu_paths, cpu_mean, 1); let cpu_cov4 = lag_covariance(&cpu_paths, cpu_mean, 4); - let cuda_result = fgn - .sample_cuda_native(m) + let cuda_batch = fgn + .sample_cuda_native_impl(m) .expect("cuda batch should succeed"); - let cuda_batch = cuda_result.right().expect("batch"); let cuda_paths: Vec> = cuda_batch.rows().into_iter().map(|r| r.to_vec()).collect(); let cuda_vals: Vec = cuda_paths.iter().flatten().copied().collect(); let cuda_mean = cuda_vals.iter().sum::() / cuda_vals.len() as f64; diff --git a/stochastic-rs-stochastic/src/noise/fgn/cuda_oxide.rs b/stochastic-rs-stochastic/src/noise/fgn/cuda_oxide.rs new file mode 100644 index 00000000..0c57cdb2 --- /dev/null +++ b/stochastic-rs-stochastic/src/noise/fgn/cuda_oxide.rs @@ -0,0 +1,294 @@ +//! # cuda-oxide Experimental CUDA +//! +//! Experimental Fgn sampling whose device kernels are written in pure Rust and +//! compiled to NVVM IR by the NVlabs `cuda-oxide` rustc codegen backend. The +//! kernels live in the device-only `fgn-oxide-kernels` crate; their NVVM IR is +//! generated once (by the maintainer, via `cargo oxide`) and committed here as +//! `fgn_oxide_kernels.ll`. It is `include_bytes!`-embedded and lowered to a +//! cubin at runtime with libNVVM + nvJitLink, so a plain downstream +//! `cargo build` runs with no `cargo oxide` / precompile step. +//! +//! This backend is intentionally separate from the stable `cuda-native` +//! backend while the cuda-oxide compiler/runtime are still alpha. + +use std::any::TypeId; +use std::sync::OnceLock; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use anyhow::Result; +use cuda_core::DeviceBuffer; +use cuda_core::LaunchConfig; +use cuda_host::cuda_launch; +use fgn_oxide_kernels::*; +use ndarray::Array2; +use stochastic_rs_core::simd_rng::SeedExt; + +use super::Fgn; +use crate::traits::FloatExt; + +/// NVVM IR for the fGN device kernels, generated once with the cuda-oxide +/// rustc codegen backend (`fgn-oxide-kernels` crate) and committed. Embedded so +/// downstream needs no `cargo oxide`; lowered to a cubin on first use below. +const FGN_OXIDE_NVVM_IR: &[u8] = include_bytes!("fgn_oxide_kernels.ll"); + +static RNG_SEQ: AtomicU64 = AtomicU64::new(0); + +/// Path to the cubin lowered from [`FGN_OXIDE_NVVM_IR`], cached for the process. +/// Compilation is deterministic per GPU arch, so one compile serves every call. +static FGN_OXIDE_CUBIN_PATH: OnceLock = OnceLock::new(); + +/// Lower the embedded NVVM IR to a cubin (libNVVM + nvJitLink, linking +/// `libdevice` for the `ln`/`sqrt`/`cos`/`sin` the RNG uses) and return its +/// path, cached after the first call. The IR is written to a temp file because +/// the runtime's `build_cubin_from_ll` is path-based. Needs the CUDA Toolkit's +/// libNVVM at runtime — but no `cargo oxide` / precompile step downstream. +fn fgn_oxide_module_path(arch: &str) -> Result<&'static str> { + if let Some(path) = FGN_OXIDE_CUBIN_PATH.get() { + return Ok(path.to_str().expect("cubin path is valid UTF-8")); + } + let dir = std::env::temp_dir().join("stochastic_rs_cuda_oxide"); + std::fs::create_dir_all(&dir).map_err(|e| anyhow::anyhow!("cuda-oxide tmp dir: {e}"))?; + let ll_path = dir.join("fgn_oxide_kernels.ll"); + std::fs::write(&ll_path, FGN_OXIDE_NVVM_IR) + .map_err(|e| anyhow::anyhow!("cuda-oxide write IR: {e}"))?; + let cubin = cuda_host::ltoir::build_cubin_from_ll(&ll_path, arch) + .map_err(|e| anyhow::anyhow!("cuda-oxide: build cubin from embedded NVVM IR: {e}"))?; + let _ = FGN_OXIDE_CUBIN_PATH.set(cubin); + Ok( + FGN_OXIDE_CUBIN_PATH + .get() + .expect("just set") + .to_str() + .expect("cubin path is valid UTF-8"), + ) +} + +fn array2_from_flat>( + host: &[U], + m: usize, + cols: usize, +) -> Array2 { + let mut out = Array2::::zeros((m, cols)); + for i in 0..m { + for j in 0..cols { + out[[i, j]] = T::from_f64_fast(host[i * cols + j].into()); + } + } + out +} + +fn array2_from_vec_f32(v: Vec, m: usize, cols: usize) -> Array2 { + if TypeId::of::() == TypeId::of::() { + let out = Array2::::from_shape_vec((m, cols), v).expect("shape must be valid"); + unsafe { std::mem::transmute::, Array2>(out) } + } else { + array2_from_flat::(&v, m, cols) + } +} + +fn array2_from_vec_f64(v: Vec, m: usize, cols: usize) -> Array2 { + if TypeId::of::() == TypeId::of::() { + let out = Array2::::from_shape_vec((m, cols), v).expect("shape must be valid"); + unsafe { std::mem::transmute::, Array2>(out) } + } else { + array2_from_flat::(&v, m, cols) + } +} + +fn sample_f32( + sqrt_eigs: &[f32], + n: usize, + m: usize, + offset: usize, + hurst: f64, + t: f64, +) -> Result> { + let out_size = n - offset; + let traj_size = 2 * n; + let total_complex = m * traj_size; + let scale = (out_size.max(1) as f32).powf(-(hurst as f32)) * (t as f32).powf(hurst as f32); + + let ctx = cuda_core::CudaContext::new(0).map_err(|e| anyhow::anyhow!("CudaContext: {e}"))?; + let stream = ctx.default_stream(); + let (cc_major, cc_minor) = ctx + .compute_capability() + .map_err(|e| anyhow::anyhow!("cuda-oxide compute capability: {e}"))?; + let arch = format!("sm_{cc_major}{cc_minor}"); + let module = ctx + .load_module_from_file(fgn_oxide_module_path(&arch)?) + .map_err(|e| anyhow::anyhow!("cuda-oxide load module: {e}"))?; + + let mut data_dev = DeviceBuffer::::zeroed(&stream, 2 * total_complex) + .map_err(|e| anyhow::anyhow!("cuda-oxide alloc complex data: {e}"))?; + let eigs_dev = DeviceBuffer::from_host(&stream, sqrt_eigs) + .map_err(|e| anyhow::anyhow!("cuda-oxide htod eigs: {e}"))?; + let mut out_dev = DeviceBuffer::::zeroed(&stream, m * out_size) + .map_err(|e| anyhow::anyhow!("cuda-oxide alloc out: {e}"))?; + + let seed: u64 = rand::Rng::random(&mut crate::simd_rng::rng()); + let seq = RNG_SEQ.fetch_add(total_complex as u64, Ordering::Relaxed); + cuda_launch! { + kernel: gen_scale_f32, + stream: stream, + module: module, + config: LaunchConfig::for_num_elems(total_complex as u32), + args: [slice_mut(data_dev), slice(eigs_dev), traj_size, total_complex, seed, seq] + } + .map_err(|e| anyhow::anyhow!("cuda-oxide gen_scale_f32: {e}"))?; + + let log_n = traj_size.trailing_zeros() as usize; + cuda_launch! { + kernel: bit_reverse_f32, + stream: stream, + module: module, + config: LaunchConfig::for_num_elems(total_complex as u32), + args: [slice_mut(data_dev), traj_size, log_n] + } + .map_err(|e| anyhow::anyhow!("cuda-oxide bit_reverse_f32: {e}"))?; + + let total_butterflies = m * (traj_size / 2); + let mut half_stride = 1usize; + while half_stride < traj_size { + cuda_launch! { + kernel: fft_stage_f32, + stream: stream, + module: module, + config: LaunchConfig::for_num_elems(total_butterflies as u32), + args: [slice_mut(data_dev), traj_size, half_stride] + } + .map_err(|e| anyhow::anyhow!("cuda-oxide fft_stage_f32({half_stride}): {e}"))?; + half_stride *= 2; + } + + cuda_launch! { + kernel: extract_real_f32, + stream: stream, + module: module, + config: LaunchConfig::for_num_elems((m * out_size) as u32), + args: [slice(data_dev), slice_mut(out_dev), out_size, traj_size, scale] + } + .map_err(|e| anyhow::anyhow!("cuda-oxide extract_real_f32: {e}"))?; + + let host = out_dev + .to_host_vec(&stream) + .map_err(|e| anyhow::anyhow!("cuda-oxide dtoh: {e}"))?; + let fgn = array2_from_vec_f32::(host, m, out_size); + Ok(fgn) +} + +fn sample_f64( + sqrt_eigs: &[f64], + n: usize, + m: usize, + offset: usize, + hurst: f64, + t: f64, +) -> Result> { + let out_size = n - offset; + let traj_size = 2 * n; + let total_complex = m * traj_size; + let scale = (out_size.max(1) as f64).powf(-hurst) * t.powf(hurst); + + let ctx = cuda_core::CudaContext::new(0).map_err(|e| anyhow::anyhow!("CudaContext: {e}"))?; + let stream = ctx.default_stream(); + let (cc_major, cc_minor) = ctx + .compute_capability() + .map_err(|e| anyhow::anyhow!("cuda-oxide compute capability: {e}"))?; + let arch = format!("sm_{cc_major}{cc_minor}"); + let module = ctx + .load_module_from_file(fgn_oxide_module_path(&arch)?) + .map_err(|e| anyhow::anyhow!("cuda-oxide load module: {e}"))?; + + let mut data_dev = DeviceBuffer::::zeroed(&stream, 2 * total_complex) + .map_err(|e| anyhow::anyhow!("cuda-oxide alloc complex data: {e}"))?; + let eigs_dev = DeviceBuffer::from_host(&stream, sqrt_eigs) + .map_err(|e| anyhow::anyhow!("cuda-oxide htod eigs: {e}"))?; + let mut out_dev = DeviceBuffer::::zeroed(&stream, m * out_size) + .map_err(|e| anyhow::anyhow!("cuda-oxide alloc out: {e}"))?; + + let seed: u64 = rand::Rng::random(&mut crate::simd_rng::rng()); + let seq = RNG_SEQ.fetch_add(total_complex as u64, Ordering::Relaxed); + cuda_launch! { + kernel: gen_scale_f64, + stream: stream, + module: module, + config: LaunchConfig::for_num_elems(total_complex as u32), + args: [slice_mut(data_dev), slice(eigs_dev), traj_size, total_complex, seed, seq] + } + .map_err(|e| anyhow::anyhow!("cuda-oxide gen_scale_f64: {e}"))?; + + let log_n = traj_size.trailing_zeros() as usize; + cuda_launch! { + kernel: bit_reverse_f64, + stream: stream, + module: module, + config: LaunchConfig::for_num_elems(total_complex as u32), + args: [slice_mut(data_dev), traj_size, log_n] + } + .map_err(|e| anyhow::anyhow!("cuda-oxide bit_reverse_f64: {e}"))?; + + let total_butterflies = m * (traj_size / 2); + let mut half_stride = 1usize; + while half_stride < traj_size { + cuda_launch! { + kernel: fft_stage_f64, + stream: stream, + module: module, + config: LaunchConfig::for_num_elems(total_butterflies as u32), + args: [slice_mut(data_dev), traj_size, half_stride] + } + .map_err(|e| anyhow::anyhow!("cuda-oxide fft_stage_f64({half_stride}): {e}"))?; + half_stride *= 2; + } + + cuda_launch! { + kernel: extract_real_f64, + stream: stream, + module: module, + config: LaunchConfig::for_num_elems((m * out_size) as u32), + args: [slice(data_dev), slice_mut(out_dev), out_size, traj_size, scale] + } + .map_err(|e| anyhow::anyhow!("cuda-oxide extract_real_f64: {e}"))?; + + let host = out_dev + .to_host_vec(&stream) + .map_err(|e| anyhow::anyhow!("cuda-oxide dtoh: {e}"))?; + let fgn = array2_from_vec_f64::(host, m, out_size); + Ok(fgn) +} + +impl Fgn { + /// Sample with the experimental cuda-oxide backend. + /// + /// Device kernels are embedded as NVVM IR and lowered to a cubin at runtime + /// (libNVVM + nvJitLink), so no `cargo oxide` precompile step is needed + /// downstream. `_module_name` is accepted for backwards compatibility but + /// ignored — the embedded module is always used. + pub fn sample_cuda_oxide_with_module(&self, m: usize, _module_name: &str) -> Result> { + let n = self.n; + let offset = self.offset; + let hurst = self.hurst.to_f64().unwrap(); + let t = self.t.unwrap_or(T::one()).to_f64().unwrap(); + + if TypeId::of::() == TypeId::of::() { + let eigs: Vec = self + .sqrt_eigenvalues + .iter() + .map(|x| x.to_f32().unwrap()) + .collect(); + return sample_f32::(&eigs, n, m, offset, hurst, t); + } + + let eigs: Vec = self + .sqrt_eigenvalues + .iter() + .map(|x| x.to_f64().unwrap()) + .collect(); + sample_f64::(&eigs, n, m, offset, hurst, t) + } + + pub(crate) fn sample_cuda_oxide_impl(&self, m: usize) -> Result> { + self.sample_cuda_oxide_with_module(m, "") + } +} diff --git a/stochastic-rs-stochastic/src/noise/fgn/fgn_oxide_kernels.ll b/stochastic-rs-stochastic/src/noise/fgn/fgn_oxide_kernels.ll new file mode 100644 index 00000000..32fc5c21 --- /dev/null +++ b/stochastic-rs-stochastic/src/noise/fgn/fgn_oxide_kernels.ll @@ -0,0 +1,741 @@ +; ModuleID = 'builtin.module' +source_filename = "fgn_oxide_kernels" +target datalayout = "e-p:64:64:64-p3:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-i128:128:128-f32:32:32-f64:64:64-f128:128:128-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64-a:8:8" +target triple = "nvptx64-nvidia-cuda" + +declare i32 @llvm.nvvm.read.ptx.sreg.tid.x() +declare i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() +declare i32 @llvm.nvvm.read.ptx.sreg.ntid.x() + +define void @bit_reverse_f32(ptr %v0, i64 %v1, i64 %v2, i64 %v3) { +entry: + %v4 = insertvalue { ptr, i64 } undef, ptr %v0, 0 + %v5 = insertvalue { ptr, i64 } %v4, i64 %v1, 1 + br label %bb0 +bb0: + %v6 = phi { ptr, i64 } [ %v5, %entry ] + %v7 = phi i64 [ %v2, %entry ] + %v8 = phi i64 [ %v3, %entry ] + %v9 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x() + br label %bb5 +bb1: + %v10 = udiv i64 %v32, %v7 + %v11 = urem i64 %v32, %v7 + br label %bb8 +bb2: + %v12 = mul i64 %v10, %v7 + %v13 = add i64 %v12, %v11 + %v14 = mul i64 2, %v13 + %v15 = add i64 %v12, %v35 + %v16 = mul i64 2, %v15 + %v17 = extractvalue { ptr, i64 } %v6, 0 + %v18 = getelementptr inbounds float, ptr %v17, i64 %v14 + %v19 = load float, ptr %v18 + %v20 = add i64 %v14, 1 + %v21 = getelementptr inbounds float, ptr %v17, i64 %v20 + %v22 = load float, ptr %v21 + %v23 = getelementptr inbounds float, ptr %v17, i64 %v16 + %v24 = load float, ptr %v23 + store float %v24, ptr %v18 + %v25 = add i64 %v16, 1 + %v26 = getelementptr inbounds float, ptr %v17, i64 %v25 + %v27 = load float, ptr %v26 + store float %v27, ptr %v21 + store float %v19, ptr %v23 + store float %v22, ptr %v26 + br label %bb4 +bb3: + br label %bb4 +bb4: + ret void +bb5: + %v28 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() + br label %bb6 +bb6: + %v29 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x() + br label %bb7 +bb7: + %v30 = mul i32 %v28, %v29 + %v31 = add i32 %v30, %v9 + %v32 = zext i32 %v31 to i64 + %v33 = icmp eq i64 %v7, 0 + %v34 = xor i1 %v33, 1 + br i1 %v34, label %bb1, label %bb11 +bb8: + %v35 = phi i64 [ 0, %bb1 ], [ %v44, %bb9 ] + %v36 = phi i64 [ %v11, %bb1 ], [ %v47, %bb9 ] + %v37 = phi i64 [ 0, %bb1 ], [ %v48, %bb9 ] + %v38 = icmp ult i64 %v37, %v8 + %v39 = xor i1 %v38, 1 + br i1 %v39, label %bb10, label %bb9 +bb9: + %v40 = zext i32 1 to i64 + %v41 = and i64 %v40, 63 + %v42 = shl i64 %v35, %v41 + %v43 = and i64 %v36, 1 + %v44 = or i64 %v42, %v43 + %v45 = zext i32 1 to i64 + %v46 = and i64 %v45, 63 + %v47 = lshr i64 %v36, %v46 + %v48 = add i64 %v37, 1 + br label %bb8 +bb10: + %v49 = icmp ult i64 %v11, %v35 + %v50 = xor i1 %v49, 1 + br i1 %v50, label %bb3, label %bb2 +bb11: + unreachable +} + +declare double @__nv_log(double) +declare double @__nv_cos(double) +declare double @__nv_sqrt(double) +declare double @__nv_sin(double) + +define void @gen_scale_f64(ptr %v0, i64 %v1, ptr %v2, i64 %v3, i64 %v4, i64 %v5, i64 %v6, i64 %v7) { +entry: + %v8 = insertvalue { ptr, i64 } undef, ptr %v0, 0 + %v9 = insertvalue { ptr, i64 } %v8, i64 %v1, 1 + %v10 = insertvalue { ptr, i64 } undef, ptr %v2, 0 + %v11 = insertvalue { ptr, i64 } %v10, i64 %v3, 1 + br label %bb0 +bb0: + %v12 = phi { ptr, i64 } [ %v9, %entry ] + %v13 = phi { ptr, i64 } [ %v11, %entry ] + %v14 = phi i64 [ %v4, %entry ] + %v15 = phi i64 [ %v5, %entry ] + %v16 = phi i64 [ %v6, %entry ] + %v17 = phi i64 [ %v7, %entry ] + %v18 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x() + br label %bb6 +bb1: + %v19 = call { i32, i32 } @fgn_oxide_kernels__philox2x32_10(i64 %v42, i64 %v16, i64 %v17) + br label %bb2 +bb2: + %v20 = extractvalue { i32, i32 } %v19, 0 + %v21 = extractvalue { i32, i32 } %v19, 1 + %v22 = uitofp i32 %v20 to double + %v23 = fadd double %v22, 0.5 + %v24 = fmul double %v23, 0.00000000023283064365386963 + %v25 = uitofp i32 %v21 to double + %v26 = fadd double %v25, 0.5 + %v27 = fmul double %v26, 0.00000000023283064365386963 + %v28 = call double @__nv_log(double %v24) + br label %bb9 +bb3: + %v29 = urem i64 %v42, %v14 + %v30 = extractvalue { ptr, i64 } %v13, 1 + %v31 = icmp ult i64 %v29, %v30 + br i1 %v31, label %bb4, label %bb13 +bb4: + %v32 = extractvalue { ptr, i64 } %v13, 0 + %v33 = getelementptr inbounds double, ptr %v32, i64 %v29 + %v34 = load double, ptr %v33 + %v35 = mul i64 2, %v42 + %v36 = extractvalue { ptr, i64 } %v12, 0 + %v37 = call double @__nv_cos(double %v47) + br label %bb11 +bb5: + ret void +bb6: + %v38 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() + br label %bb7 +bb7: + %v39 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x() + br label %bb8 +bb8: + %v40 = mul i32 %v38, %v39 + %v41 = add i32 %v40, %v18 + %v42 = zext i32 %v41 to i64 + %v43 = icmp ult i64 %v42, %v15 + %v44 = xor i1 %v43, 1 + br i1 %v44, label %bb5, label %bb1 +bb9: + %v45 = fmul double -2.0, %v28 + %v46 = call double @__nv_sqrt(double %v45) + br label %bb10 +bb10: + %v47 = fmul double 6.283185307179586, %v27 + %v48 = icmp eq i64 %v14, 0 + %v49 = xor i1 %v48, 1 + br i1 %v49, label %bb3, label %bb14 +bb11: + %v50 = fmul double %v46, %v37 + %v51 = getelementptr inbounds double, ptr %v36, i64 %v35 + %v52 = fmul double %v50, %v34 + store double %v52, ptr %v51 + %v53 = call double @__nv_sin(double %v47) + br label %bb12 +bb12: + %v54 = fmul double %v46, %v53 + %v55 = add i64 %v35, 1 + %v56 = getelementptr inbounds double, ptr %v36, i64 %v55 + %v57 = fmul double %v54, %v34 + store double %v57, ptr %v56 + br label %bb5 +bb13: + unreachable +bb14: + unreachable +} + +define void @bit_reverse_f64(ptr %v0, i64 %v1, i64 %v2, i64 %v3) { +entry: + %v4 = insertvalue { ptr, i64 } undef, ptr %v0, 0 + %v5 = insertvalue { ptr, i64 } %v4, i64 %v1, 1 + br label %bb0 +bb0: + %v6 = phi { ptr, i64 } [ %v5, %entry ] + %v7 = phi i64 [ %v2, %entry ] + %v8 = phi i64 [ %v3, %entry ] + %v9 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x() + br label %bb5 +bb1: + %v10 = udiv i64 %v32, %v7 + %v11 = urem i64 %v32, %v7 + br label %bb8 +bb2: + %v12 = mul i64 %v10, %v7 + %v13 = add i64 %v12, %v11 + %v14 = mul i64 2, %v13 + %v15 = add i64 %v12, %v35 + %v16 = mul i64 2, %v15 + %v17 = extractvalue { ptr, i64 } %v6, 0 + %v18 = getelementptr inbounds double, ptr %v17, i64 %v14 + %v19 = load double, ptr %v18 + %v20 = add i64 %v14, 1 + %v21 = getelementptr inbounds double, ptr %v17, i64 %v20 + %v22 = load double, ptr %v21 + %v23 = getelementptr inbounds double, ptr %v17, i64 %v16 + %v24 = load double, ptr %v23 + store double %v24, ptr %v18 + %v25 = add i64 %v16, 1 + %v26 = getelementptr inbounds double, ptr %v17, i64 %v25 + %v27 = load double, ptr %v26 + store double %v27, ptr %v21 + store double %v19, ptr %v23 + store double %v22, ptr %v26 + br label %bb4 +bb3: + br label %bb4 +bb4: + ret void +bb5: + %v28 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() + br label %bb6 +bb6: + %v29 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x() + br label %bb7 +bb7: + %v30 = mul i32 %v28, %v29 + %v31 = add i32 %v30, %v9 + %v32 = zext i32 %v31 to i64 + %v33 = icmp eq i64 %v7, 0 + %v34 = xor i1 %v33, 1 + br i1 %v34, label %bb1, label %bb11 +bb8: + %v35 = phi i64 [ 0, %bb1 ], [ %v44, %bb9 ] + %v36 = phi i64 [ %v11, %bb1 ], [ %v47, %bb9 ] + %v37 = phi i64 [ 0, %bb1 ], [ %v48, %bb9 ] + %v38 = icmp ult i64 %v37, %v8 + %v39 = xor i1 %v38, 1 + br i1 %v39, label %bb10, label %bb9 +bb9: + %v40 = zext i32 1 to i64 + %v41 = and i64 %v40, 63 + %v42 = shl i64 %v35, %v41 + %v43 = and i64 %v36, 1 + %v44 = or i64 %v42, %v43 + %v45 = zext i32 1 to i64 + %v46 = and i64 %v45, 63 + %v47 = lshr i64 %v36, %v46 + %v48 = add i64 %v37, 1 + br label %bb8 +bb10: + %v49 = icmp ult i64 %v11, %v35 + %v50 = xor i1 %v49, 1 + br i1 %v50, label %bb3, label %bb2 +bb11: + unreachable +} + +define void @extract_real_f64(ptr %v0, i64 %v1, ptr %v2, i64 %v3, i64 %v4, i64 %v5, double %v6) { +entry: + %v7 = insertvalue { ptr, i64 } undef, ptr %v0, 0 + %v8 = insertvalue { ptr, i64 } %v7, i64 %v1, 1 + %v9 = insertvalue { ptr, i64 } undef, ptr %v2, 0 + %v10 = insertvalue { ptr, i64 } %v9, i64 %v3, 1 + br label %bb0 +bb0: + %v11 = phi { ptr, i64 } [ %v8, %entry ] + %v12 = phi { ptr, i64 } [ %v10, %entry ] + %v13 = phi i64 [ %v4, %entry ] + %v14 = phi i64 [ %v5, %entry ] + %v15 = phi double [ %v6, %entry ] + %v16 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x() + br label %bb6 +bb1: + %v17 = extractvalue { i8, ptr } %v45, 1 + %v18 = icmp eq i64 %v13, 0 + %v19 = xor i1 %v18, 1 + br i1 %v19, label %bb2, label %bb14 +bb2: + %v20 = udiv i64 %v36, %v13 + %v21 = urem i64 %v36, %v13 + %v22 = mul i64 %v20, %v14 + %v23 = add i64 %v22, %v21 + %v24 = add i64 %v23, 1 + %v25 = mul i64 2, %v24 + %v26 = extractvalue { ptr, i64 } %v11, 1 + %v27 = icmp ult i64 %v25, %v26 + br i1 %v27, label %bb3, label %bb15 +bb3: + %v28 = extractvalue { ptr, i64 } %v11, 0 + %v29 = getelementptr inbounds double, ptr %v28, i64 %v25 + %v30 = load double, ptr %v29 + %v31 = fmul double %v30, %v15 + store double %v31, ptr %v17 + br label %bb5 +bb4: + br label %bb5 +bb5: + ret void +bb6: + %v32 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() + br label %bb7 +bb7: + %v33 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x() + br label %bb8 +bb8: + %v34 = mul i32 %v32, %v33 + %v35 = add i32 %v34, %v16 + %v36 = zext i32 %v35 to i64 + %v37 = extractvalue { ptr, i64 } %v12, 1 + %v38 = icmp ult i64 %v36, %v37 + %v39 = xor i1 %v38, 1 + br i1 %v39, label %bb10, label %bb9 +bb9: + %v40 = extractvalue { ptr, i64 } %v12, 0 + %v41 = getelementptr inbounds double, ptr %v40, i64 %v36 + %v42 = insertvalue { i8, ptr } undef, i8 1, 0 + %v43 = insertvalue { i8, ptr } %v42, ptr %v41, 1 + br label %bb11 +bb10: + %v44 = insertvalue { i8, ptr } undef, i8 0, 0 + br label %bb11 +bb11: + %v45 = phi { i8, ptr } [ %v43, %bb9 ], [ %v44, %bb10 ] + %v46 = extractvalue { i8, ptr } %v45, 0 + %v47 = zext i8 %v46 to i64 + %v48 = icmp eq i64 %v47, 1 + br i1 %v48, label %bb1, label %bb12 +bb12: + %v49 = icmp eq i64 %v47, 0 + br i1 %v49, label %bb4, label %bb13 +bb13: + unreachable +bb14: + unreachable +bb15: + unreachable +} + +define void @extract_real_f32(ptr %v0, i64 %v1, ptr %v2, i64 %v3, i64 %v4, i64 %v5, float %v6) { +entry: + %v7 = insertvalue { ptr, i64 } undef, ptr %v0, 0 + %v8 = insertvalue { ptr, i64 } %v7, i64 %v1, 1 + %v9 = insertvalue { ptr, i64 } undef, ptr %v2, 0 + %v10 = insertvalue { ptr, i64 } %v9, i64 %v3, 1 + br label %bb0 +bb0: + %v11 = phi { ptr, i64 } [ %v8, %entry ] + %v12 = phi { ptr, i64 } [ %v10, %entry ] + %v13 = phi i64 [ %v4, %entry ] + %v14 = phi i64 [ %v5, %entry ] + %v15 = phi float [ %v6, %entry ] + %v16 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x() + br label %bb6 +bb1: + %v17 = extractvalue { i8, ptr } %v45, 1 + %v18 = icmp eq i64 %v13, 0 + %v19 = xor i1 %v18, 1 + br i1 %v19, label %bb2, label %bb14 +bb2: + %v20 = udiv i64 %v36, %v13 + %v21 = urem i64 %v36, %v13 + %v22 = mul i64 %v20, %v14 + %v23 = add i64 %v22, %v21 + %v24 = add i64 %v23, 1 + %v25 = mul i64 2, %v24 + %v26 = extractvalue { ptr, i64 } %v11, 1 + %v27 = icmp ult i64 %v25, %v26 + br i1 %v27, label %bb3, label %bb15 +bb3: + %v28 = extractvalue { ptr, i64 } %v11, 0 + %v29 = getelementptr inbounds float, ptr %v28, i64 %v25 + %v30 = load float, ptr %v29 + %v31 = fmul float %v30, %v15 + store float %v31, ptr %v17 + br label %bb5 +bb4: + br label %bb5 +bb5: + ret void +bb6: + %v32 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() + br label %bb7 +bb7: + %v33 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x() + br label %bb8 +bb8: + %v34 = mul i32 %v32, %v33 + %v35 = add i32 %v34, %v16 + %v36 = zext i32 %v35 to i64 + %v37 = extractvalue { ptr, i64 } %v12, 1 + %v38 = icmp ult i64 %v36, %v37 + %v39 = xor i1 %v38, 1 + br i1 %v39, label %bb10, label %bb9 +bb9: + %v40 = extractvalue { ptr, i64 } %v12, 0 + %v41 = getelementptr inbounds float, ptr %v40, i64 %v36 + %v42 = insertvalue { i8, ptr } undef, i8 1, 0 + %v43 = insertvalue { i8, ptr } %v42, ptr %v41, 1 + br label %bb11 +bb10: + %v44 = insertvalue { i8, ptr } undef, i8 0, 0 + br label %bb11 +bb11: + %v45 = phi { i8, ptr } [ %v43, %bb9 ], [ %v44, %bb10 ] + %v46 = extractvalue { i8, ptr } %v45, 0 + %v47 = zext i8 %v46 to i64 + %v48 = icmp eq i64 %v47, 1 + br i1 %v48, label %bb1, label %bb12 +bb12: + %v49 = icmp eq i64 %v47, 0 + br i1 %v49, label %bb4, label %bb13 +bb13: + unreachable +bb14: + unreachable +bb15: + unreachable +} + +declare float @__nv_logf(float) +declare float @__nv_cosf(float) +declare float @__nv_sqrtf(float) +declare float @__nv_sinf(float) + +define void @gen_scale_f32(ptr %v0, i64 %v1, ptr %v2, i64 %v3, i64 %v4, i64 %v5, i64 %v6, i64 %v7) { +entry: + %v8 = insertvalue { ptr, i64 } undef, ptr %v0, 0 + %v9 = insertvalue { ptr, i64 } %v8, i64 %v1, 1 + %v10 = insertvalue { ptr, i64 } undef, ptr %v2, 0 + %v11 = insertvalue { ptr, i64 } %v10, i64 %v3, 1 + br label %bb0 +bb0: + %v12 = phi { ptr, i64 } [ %v9, %entry ] + %v13 = phi { ptr, i64 } [ %v11, %entry ] + %v14 = phi i64 [ %v4, %entry ] + %v15 = phi i64 [ %v5, %entry ] + %v16 = phi i64 [ %v6, %entry ] + %v17 = phi i64 [ %v7, %entry ] + %v18 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x() + br label %bb6 +bb1: + %v19 = call { i32, i32 } @fgn_oxide_kernels__philox2x32_10(i64 %v42, i64 %v16, i64 %v17) + br label %bb2 +bb2: + %v20 = extractvalue { i32, i32 } %v19, 0 + %v21 = extractvalue { i32, i32 } %v19, 1 + %v22 = uitofp i32 %v20 to float + %v23 = fadd float %v22, 0.5 + %v24 = fmul float %v23, 0.00000000023283064365386963 + %v25 = uitofp i32 %v21 to float + %v26 = fadd float %v25, 0.5 + %v27 = fmul float %v26, 0.00000000023283064365386963 + %v28 = call float @__nv_logf(float %v24) + br label %bb9 +bb3: + %v29 = urem i64 %v42, %v14 + %v30 = extractvalue { ptr, i64 } %v13, 1 + %v31 = icmp ult i64 %v29, %v30 + br i1 %v31, label %bb4, label %bb13 +bb4: + %v32 = extractvalue { ptr, i64 } %v13, 0 + %v33 = getelementptr inbounds float, ptr %v32, i64 %v29 + %v34 = load float, ptr %v33 + %v35 = mul i64 2, %v42 + %v36 = extractvalue { ptr, i64 } %v12, 0 + %v37 = call float @__nv_cosf(float %v47) + br label %bb11 +bb5: + ret void +bb6: + %v38 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() + br label %bb7 +bb7: + %v39 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x() + br label %bb8 +bb8: + %v40 = mul i32 %v38, %v39 + %v41 = add i32 %v40, %v18 + %v42 = zext i32 %v41 to i64 + %v43 = icmp ult i64 %v42, %v15 + %v44 = xor i1 %v43, 1 + br i1 %v44, label %bb5, label %bb1 +bb9: + %v45 = fmul float -2.0, %v28 + %v46 = call float @__nv_sqrtf(float %v45) + br label %bb10 +bb10: + %v47 = fmul float 6.2831854820251465, %v27 + %v48 = icmp eq i64 %v14, 0 + %v49 = xor i1 %v48, 1 + br i1 %v49, label %bb3, label %bb14 +bb11: + %v50 = fmul float %v46, %v37 + %v51 = getelementptr inbounds float, ptr %v36, i64 %v35 + %v52 = fmul float %v50, %v34 + store float %v52, ptr %v51 + %v53 = call float @__nv_sinf(float %v47) + br label %bb12 +bb12: + %v54 = fmul float %v46, %v53 + %v55 = add i64 %v35, 1 + %v56 = getelementptr inbounds float, ptr %v36, i64 %v55 + %v57 = fmul float %v54, %v34 + store float %v57, ptr %v56 + br label %bb5 +bb13: + unreachable +bb14: + unreachable +} + +define void @fft_stage_f64(ptr %v0, i64 %v1, i64 %v2, i64 %v3) { +entry: + %v4 = insertvalue { ptr, i64 } undef, ptr %v0, 0 + %v5 = insertvalue { ptr, i64 } %v4, i64 %v1, 1 + br label %bb0 +bb0: + %v6 = phi { ptr, i64 } [ %v5, %entry ] + %v7 = phi i64 [ %v2, %entry ] + %v8 = phi i64 [ %v3, %entry ] + %v9 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x() + br label %bb3 +bb1: + %v10 = udiv i64 %v31, %v32 + %v11 = urem i64 %v31, %v32 + %v12 = mul i64 %v8, 2 + %v13 = icmp eq i64 %v8, 0 + %v14 = xor i1 %v13, 1 + br i1 %v14, label %bb2, label %bb8 +bb2: + %v15 = udiv i64 %v11, %v8 + %v16 = urem i64 %v11, %v8 + %v17 = mul i64 %v10, %v7 + %v18 = mul i64 %v15, %v12 + %v19 = add i64 %v17, %v18 + %v20 = add i64 %v19, %v16 + %v21 = add i64 %v20, %v8 + %v22 = uitofp i64 %v16 to double + %v23 = fmul double -6.283185307179586, %v22 + %v24 = uitofp i64 %v12 to double + %v25 = fdiv double %v23, %v24 + %v26 = call double @__nv_cos(double %v25) + br label %bb6 +bb3: + %v27 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() + br label %bb4 +bb4: + %v28 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x() + br label %bb5 +bb5: + %v29 = mul i32 %v27, %v28 + %v30 = add i32 %v29, %v9 + %v31 = zext i32 %v30 to i64 + %v32 = udiv i64 %v7, 2 + %v33 = icmp eq i64 %v32, 0 + %v34 = xor i1 %v33, 1 + br i1 %v34, label %bb1, label %bb9 +bb6: + %v35 = call double @__nv_sin(double %v25) + br label %bb7 +bb7: + %v36 = extractvalue { ptr, i64 } %v6, 0 + %v37 = mul i64 2, %v20 + %v38 = mul i64 2, %v21 + %v39 = getelementptr inbounds double, ptr %v36, i64 %v37 + %v40 = load double, ptr %v39 + %v41 = add i64 %v37, 1 + %v42 = getelementptr inbounds double, ptr %v36, i64 %v41 + %v43 = load double, ptr %v42 + %v44 = getelementptr inbounds double, ptr %v36, i64 %v38 + %v45 = load double, ptr %v44 + %v46 = add i64 %v38, 1 + %v47 = getelementptr inbounds double, ptr %v36, i64 %v46 + %v48 = load double, ptr %v47 + %v49 = fmul double %v45, %v26 + %v50 = fmul double %v48, %v35 + %v51 = fsub double %v49, %v50 + %v52 = fmul double %v45, %v35 + %v53 = fmul double %v48, %v26 + %v54 = fadd double %v52, %v53 + %v55 = fadd double %v40, %v51 + store double %v55, ptr %v39 + %v56 = fadd double %v43, %v54 + store double %v56, ptr %v42 + %v57 = fsub double %v40, %v51 + store double %v57, ptr %v44 + %v58 = fsub double %v43, %v54 + store double %v58, ptr %v47 + ret void +bb8: + unreachable +bb9: + unreachable +} + +define void @fft_stage_f32(ptr %v0, i64 %v1, i64 %v2, i64 %v3) { +entry: + %v4 = insertvalue { ptr, i64 } undef, ptr %v0, 0 + %v5 = insertvalue { ptr, i64 } %v4, i64 %v1, 1 + br label %bb0 +bb0: + %v6 = phi { ptr, i64 } [ %v5, %entry ] + %v7 = phi i64 [ %v2, %entry ] + %v8 = phi i64 [ %v3, %entry ] + %v9 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x() + br label %bb3 +bb1: + %v10 = udiv i64 %v31, %v32 + %v11 = urem i64 %v31, %v32 + %v12 = mul i64 %v8, 2 + %v13 = icmp eq i64 %v8, 0 + %v14 = xor i1 %v13, 1 + br i1 %v14, label %bb2, label %bb8 +bb2: + %v15 = udiv i64 %v11, %v8 + %v16 = urem i64 %v11, %v8 + %v17 = mul i64 %v10, %v7 + %v18 = mul i64 %v15, %v12 + %v19 = add i64 %v17, %v18 + %v20 = add i64 %v19, %v16 + %v21 = add i64 %v20, %v8 + %v22 = uitofp i64 %v16 to float + %v23 = fmul float -6.2831854820251465, %v22 + %v24 = uitofp i64 %v12 to float + %v25 = fdiv float %v23, %v24 + %v26 = call float @__nv_cosf(float %v25) + br label %bb6 +bb3: + %v27 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() + br label %bb4 +bb4: + %v28 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x() + br label %bb5 +bb5: + %v29 = mul i32 %v27, %v28 + %v30 = add i32 %v29, %v9 + %v31 = zext i32 %v30 to i64 + %v32 = udiv i64 %v7, 2 + %v33 = icmp eq i64 %v32, 0 + %v34 = xor i1 %v33, 1 + br i1 %v34, label %bb1, label %bb9 +bb6: + %v35 = call float @__nv_sinf(float %v25) + br label %bb7 +bb7: + %v36 = extractvalue { ptr, i64 } %v6, 0 + %v37 = mul i64 2, %v20 + %v38 = mul i64 2, %v21 + %v39 = getelementptr inbounds float, ptr %v36, i64 %v37 + %v40 = load float, ptr %v39 + %v41 = add i64 %v37, 1 + %v42 = getelementptr inbounds float, ptr %v36, i64 %v41 + %v43 = load float, ptr %v42 + %v44 = getelementptr inbounds float, ptr %v36, i64 %v38 + %v45 = load float, ptr %v44 + %v46 = add i64 %v38, 1 + %v47 = getelementptr inbounds float, ptr %v36, i64 %v46 + %v48 = load float, ptr %v47 + %v49 = fmul float %v45, %v26 + %v50 = fmul float %v48, %v35 + %v51 = fsub float %v49, %v50 + %v52 = fmul float %v45, %v35 + %v53 = fmul float %v48, %v26 + %v54 = fadd float %v52, %v53 + %v55 = fadd float %v40, %v51 + store float %v55, ptr %v39 + %v56 = fadd float %v43, %v54 + store float %v56, ptr %v42 + %v57 = fsub float %v40, %v51 + store float %v57, ptr %v44 + %v58 = fsub float %v43, %v54 + store float %v58, ptr %v47 + ret void +bb8: + unreachable +bb9: + unreachable +} + +define { i32, i32 } @fgn_oxide_kernels__philox2x32_10(i64 %v0, i64 %v1, i64 %v2) { +entry: + br label %bb0 +bb0: + %v3 = phi i64 [ %v0, %entry ] + %v4 = phi i64 [ %v1, %entry ] + %v5 = phi i64 [ %v2, %entry ] + %v6 = bitcast i64 %v3 to i64 + %v7 = add i64 %v6, %v5 + %v8 = trunc i64 %v7 to i32 + %v9 = zext i32 32 to i64 + %v10 = and i64 %v9, 63 + %v11 = lshr i64 %v7, %v10 + %v12 = trunc i64 %v11 to i32 + %v13 = trunc i64 %v4 to i32 + br label %bb1 +bb1: + %v14 = phi i32 [ %v8, %bb0 ], [ %v27, %bb2 ] + %v15 = phi i32 [ %v12, %bb0 ], [ %v28, %bb2 ] + %v16 = phi i32 [ %v13, %bb0 ], [ %v29, %bb2 ] + %v17 = phi i32 [ 0, %bb0 ], [ %v30, %bb2 ] + %v18 = icmp slt i32 %v17, 10 + %v19 = xor i1 %v18, 1 + br i1 %v19, label %bb3, label %bb2 +bb2: + %v20 = zext i32 %v14 to i64 + %v21 = mul i64 3528531795, %v20 + %v22 = zext i32 32 to i64 + %v23 = and i64 %v22, 63 + %v24 = lshr i64 %v21, %v23 + %v25 = trunc i64 %v24 to i32 + %v26 = xor i32 %v25, %v15 + %v27 = xor i32 %v26, %v16 + %v28 = trunc i64 %v21 to i32 + %v29 = add i32 %v16, 2654435769 + %v30 = add i32 %v17, 1 + br label %bb1 +bb3: + %v31 = insertvalue { i32, i32 } undef, i32 %v14, 0 + %v32 = insertvalue { i32, i32 } %v31, i32 %v15, 1 + ret { i32, i32 } %v32 +} + + +@llvm.used = appending global [8 x ptr] [ptr @bit_reverse_f32, ptr @gen_scale_f64, ptr @bit_reverse_f64, ptr @extract_real_f64, ptr @extract_real_f32, ptr @gen_scale_f32, ptr @fft_stage_f64, ptr @fft_stage_f32], section "llvm.metadata" + +!0 = !{ptr @bit_reverse_f32, !"kernel", i32 1} +!1 = !{ptr @gen_scale_f64, !"kernel", i32 1} +!2 = !{ptr @bit_reverse_f64, !"kernel", i32 1} +!3 = !{ptr @extract_real_f64, !"kernel", i32 1} +!4 = !{ptr @extract_real_f32, !"kernel", i32 1} +!5 = !{ptr @gen_scale_f32, !"kernel", i32 1} +!6 = !{ptr @fft_stage_f64, !"kernel", i32 1} +!7 = !{ptr @fft_stage_f32, !"kernel", i32 1} +!nvvm.annotations = !{!0, !1, !2, !3, !4, !5, !6, !7} + +!nvvmir.version = !{!8} +!8 = !{i32 2, i32 0, i32 3, i32 2} diff --git a/stochastic-rs-stochastic/src/noise/fgn/gpu.rs b/stochastic-rs-stochastic/src/noise/fgn/gpu.rs index ed946729..106ed6b9 100644 --- a/stochastic-rs-stochastic/src/noise/fgn/gpu.rs +++ b/stochastic-rs-stochastic/src/noise/fgn/gpu.rs @@ -8,8 +8,6 @@ //! use anyhow::Result; use cubecl::prelude::*; -use either::Either; -use ndarray::Array1; use ndarray::Array2; use stochastic_rs_core::simd_rng::SeedExt; @@ -23,7 +21,12 @@ const LOCAL_STAGES: usize = 9; // log2(512) /// Shared-memory sub-FFT: loads a contiguous tile of BLOCK elements, /// performs LOCAL_STAGES radix-2 butterfly stages entirely in shared /// memory (one sync per stage), then writes back. -#[allow(clippy::approx_constant, clippy::excessive_precision)] +#[allow( + clippy::approx_constant, + clippy::excessive_precision, + clippy::identity_op, + clippy::modulo_one +)] #[cube(launch)] fn fft_local(real: &mut Array, imag: &mut Array) { let tid = UNIT_POS as usize; @@ -310,7 +313,7 @@ mod backend { hurst: f64, t: f64, seed: &S, - ) -> Result, Array2>> { + ) -> Result> { let hb = hurst.to_bits(); let tb = t.to_bits(); let traj_size = 2 * n; @@ -394,10 +397,7 @@ mod backend { let out = f32::from_bytes(&bytes); let fgn = arr2::(out, m, out_size); drop(guard); - if m == 1 { - return Ok(Either::Left(fgn.row(0).to_owned())); - } - Ok(Either::Right(fgn)) + Ok(fgn) } fn arr2(data: &[f32], m: usize, cols: usize) -> Array2 { @@ -416,8 +416,8 @@ mod backend { } } -impl Fgn { - pub(crate) fn sample_gpu_impl(&self, m: usize) -> Result, Array2>> { +impl Fgn { + pub(crate) fn sample_gpu_impl(&self, m: usize) -> Result> { #[cfg(not(any(feature = "gpu-cuda", feature = "gpu-wgpu")))] { let _ = m; diff --git a/stochastic-rs-stochastic/src/noise/fgn/metal.rs b/stochastic-rs-stochastic/src/noise/fgn/metal.rs index a00c1a74..5414f0c6 100644 --- a/stochastic-rs-stochastic/src/noise/fgn/metal.rs +++ b/stochastic-rs-stochastic/src/noise/fgn/metal.rs @@ -8,9 +8,7 @@ use std::any::TypeId; use anyhow::Result; -use either::Either; use metal::*; -use ndarray::Array1; use ndarray::Array2; use parking_lot::Mutex; use stochastic_rs_core::simd_rng::SeedExt; @@ -174,7 +172,7 @@ fn sample_f32( hurst: f64, t: f64, seed_src: &S, -) -> Result, Array2>> { +) -> Result> { let traj_size = 2 * n; let out_size = n - offset; let scale = (out_size.max(1) as f32).powf(-(hurst as f32)) * (t as f32).powf(hurst as f32); @@ -229,7 +227,7 @@ fn sample_f32( // 2. FFT butterfly stages let grid_fft = MTLSize::new((total / 2) as u64, 1, 1); for stage in 0..log_n { - let hs = (1u32 << stage) as u32; + let hs = 1u32 << stage; let enc = cmd.new_compute_command_encoder(); enc.set_compute_pipeline_state(&ctx.butterfly_pso); enc.set_buffer(0, Some(&real_buf), 0); @@ -262,10 +260,7 @@ fn sample_f32( let out_slice = unsafe { std::slice::from_raw_parts(out_ptr, m * out_size) }; let fgn = arr2_f32::(out_slice, m, out_size); - if m == 1 { - return Ok(Either::Left(fgn.row(0).to_owned())); - } - Ok(Either::Right(fgn)) + Ok(fgn) } fn arr2_f32(data: &[f32], m: usize, cols: usize) -> Array2 { @@ -283,8 +278,8 @@ fn arr2_f32(data: &[f32], m: usize, cols: usize) -> Array2 { } } -impl Fgn { - pub(crate) fn sample_metal_impl(&self, m: usize) -> Result, Array2>> { +impl Fgn { + pub(crate) fn sample_metal_impl(&self, m: usize) -> Result> { let n = self.n; let offset = self.offset; let hurst = self.hurst.to_f64().unwrap(); diff --git a/stochastic-rs-stochastic/src/process/fbm.rs b/stochastic-rs-stochastic/src/process/fbm.rs index 9bc0cad9..6c51ffd2 100644 --- a/stochastic-rs-stochastic/src/process/fbm.rs +++ b/stochastic-rs-stochastic/src/process/fbm.rs @@ -4,29 +4,10 @@ //! \mathbb E[B_t^H B_s^H]=\tfrac12\left(t^{2H}+s^{2H}-|t-s|^{2H}\right) //! $$ //! -#[cfg(any( - feature = "gpu", - feature = "cuda-native", - feature = "accelerate", - feature = "metal" -))] -use anyhow::Result; -#[cfg(any( - feature = "gpu", - feature = "cuda-native", - feature = "accelerate", - feature = "metal" -))] -use either::Either; use ndarray::Array1; -#[cfg(any( - feature = "gpu", - feature = "cuda-native", - feature = "accelerate", - feature = "metal", - feature = "python" -))] +#[cfg(feature = "python")] use ndarray::Array2; +use ndarray::parallel::prelude::*; #[cfg(feature = "python")] use numpy::IntoPyArray; #[cfg(feature = "python")] @@ -40,11 +21,13 @@ use stochastic_rs_core::simd_rng::Deterministic; use stochastic_rs_core::simd_rng::SeedExt; use stochastic_rs_core::simd_rng::Unseeded; +use crate::device::Backend; +use crate::device::Cpu; use crate::noise::fgn::Fgn; use crate::traits::FloatExt; use crate::traits::ProcessExt; -pub struct Fbm { +pub struct Fbm { /// Hurst parameter (`0 < H < 1`) controlling roughness and memory. pub hurst: T, /// Number of discrete time points in the generated path. @@ -53,10 +36,10 @@ pub struct Fbm { pub t: Option, /// Seed strategy (compile-time: [`Unseeded`] or [`Deterministic`]). pub seed: S, - fgn: Fgn, + fgn: Fgn, } -impl Fbm { +impl Fbm { pub fn new(hurst: T, n: usize, t: Option, seed: S) -> Self { assert!(n >= 2, "n must be at least 2"); @@ -70,11 +53,11 @@ impl Fbm { } } -impl ProcessExt for Fbm { +impl ProcessExt for Fbm { type Output = Array1; fn sample(&self) -> Self::Output { - let fgn = self.fgn.sample_cpu_impl(&self.seed.derive()); + let fgn = self.fgn.noise(&self.seed.derive()); let mut fbm = Array1::::zeros(self.n); for i in 1..self.n { @@ -84,61 +67,27 @@ impl ProcessExt for Fbm { fbm } - #[cfg(feature = "gpu")] - fn sample_gpu(&self, m: usize) -> Result, Array2>> { - Self::fgn_to_fbm(self.n, self.fgn.sample_gpu(m)?) - } - - #[cfg(feature = "cuda-native")] - fn sample_cuda_native(&self, m: usize) -> Result, Array2>> { - Self::fgn_to_fbm(self.n, self.fgn.sample_cuda_native(m)?) - } - - #[cfg(feature = "accelerate")] - fn sample_accelerate(&self, m: usize) -> Result, Array2>> { - Self::fgn_to_fbm(self.n, self.fgn.sample_accelerate(m)?) - } - - #[cfg(feature = "metal")] - fn sample_metal(&self, m: usize) -> Result, Array2>> { - Self::fgn_to_fbm(self.n, self.fgn.sample_metal(m)?) - } -} - -#[cfg(any( - feature = "gpu", - feature = "cuda-native", - feature = "accelerate", - feature = "metal" -))] -impl Fbm { - fn fgn_to_fbm( - n: usize, - fgn_out: Either, Array2>, - ) -> Result, Array2>> { - match fgn_out { - Either::Left(fgn_path) => { - let mut fbm = Array1::::zeros(n); - for i in 1..n { - fbm[i] = fbm[i - 1] + fgn_path[i - 1]; - } - Ok(Either::Left(fbm)) - } - Either::Right(fgn_paths) => { - let rows = fgn_paths.nrows(); - let mut fbm_paths = Array2::::zeros((rows, n)); - for r in 0..rows { - for i in 1..n { - fbm_paths[[r, i]] = fbm_paths[[r, i - 1]] + fgn_paths[[r, i - 1]]; - } + /// The `m` fGN noises are generated in one batched backend call, then each + /// path is assembled (cumulative sum) on the host across all cores. + fn sample_par(&self, m: usize) -> Vec { + self + .fgn + .noise_batch(m) + .into_par_iter() + .map(|fgn_row| { + let mut fbm = Array1::::zeros(self.n); + for i in 1..self.n { + fbm[i] = fbm[i - 1] + fgn_row[i - 1]; } - Ok(Either::Right(fbm_paths)) - } - } + fbm + }) + .collect() } } -impl Fbm { +backend_switch!([T: FloatExt, S: SeedExt] Fbm { hurst, n, t, seed } via fgn); + +impl Fbm { /// Calculate the Malliavin derivative /// /// The Malliavin derivative of the fractional Brownian motion is given by: diff --git a/stochastic-rs-stochastic/src/sde/fractional.rs b/stochastic-rs-stochastic/src/sde/fractional.rs index da200a9a..2419c9f1 100644 --- a/stochastic-rs-stochastic/src/sde/fractional.rs +++ b/stochastic-rs-stochastic/src/sde/fractional.rs @@ -7,9 +7,10 @@ use ndarray::Axis; use ndarray::s; use super::Sde; +use crate::device::Backend; use crate::traits::FloatExt; -impl Sde +impl Sde where F: Fn(&Array1, T) -> Array1, G: Fn(&Array1, T) -> Array2, diff --git a/stochastic-rs-stochastic/src/sde/gaussian.rs b/stochastic-rs-stochastic/src/sde/gaussian.rs index 288c0fe6..c2566995 100644 --- a/stochastic-rs-stochastic/src/sde/gaussian.rs +++ b/stochastic-rs-stochastic/src/sde/gaussian.rs @@ -7,9 +7,10 @@ use ndarray::s; use rand::Rng; use super::Sde; +use crate::device::Backend; use crate::traits::FloatExt; -impl Sde +impl Sde where F: Fn(&Array1, T) -> Array1, G: Fn(&Array1, T) -> ndarray::Array2, diff --git a/stochastic-rs-stochastic/src/sde/mod.rs b/stochastic-rs-stochastic/src/sde/mod.rs index 028ac283..c9c6ca94 100644 --- a/stochastic-rs-stochastic/src/sde/mod.rs +++ b/stochastic-rs-stochastic/src/sde/mod.rs @@ -179,6 +179,8 @@ mod gaussian; #[cfg(test)] mod tests; +use std::marker::PhantomData; + use ndarray::Array1; use ndarray::Array2; use ndarray::Array3; @@ -188,6 +190,8 @@ use rand::Rng; use stochastic_rs_core::simd_rng::Unseeded; use super::noise::fgn::Fgn; +use crate::device::Backend; +use crate::device::Cpu; use crate::traits::FloatExt; use crate::traits::ProcessExt; @@ -225,7 +229,7 @@ pub enum SdeMethod { /// - `G`: Diffusion function `b(x, t) -> R^{d×d}`. Takes the current state and time, /// returns the diffusion matrix where entry `[i, j]` maps noise dimension `j` to /// state dimension `i`. -pub struct Sde +pub struct Sde where F: Fn(&Array1, T) -> Array1, G: Fn(&Array1, T) -> Array2, @@ -239,9 +243,10 @@ where /// Hurst parameters for fractional noise, one per dimension. /// Required when `noise` is [`NoiseModel::Fractional`], ignored otherwise. pub hursts: Option>, + _backend: PhantomData, } -impl Sde +impl Sde where F: Fn(&Array1, T) -> Array1, G: Fn(&Array1, T) -> Array2, @@ -260,9 +265,19 @@ where diffusion, noise, hursts, + _backend: PhantomData, } } +} +backend_switch!([T: FloatExt, F, G] Sde { drift, diffusion, noise, hursts } via phantom + where F: Fn(&Array1, T) -> Array1, G: Fn(&Array1, T) -> Array2); + +impl Sde +where + F: Fn(&Array1, T) -> Array1, + G: Fn(&Array1, T) -> Array2, +{ /// Solves the SDE and returns simulated paths. /// /// # Arguments @@ -302,8 +317,8 @@ where let mut incs = Array3::zeros((n_paths, steps, dim)); if let Some(h) = &self.hursts { - let fgns: Vec> = (0..dim) - .map(|d| Fgn::new(h[d], steps, Some(t1 - t0), Unseeded)) + let fgns: Vec> = (0..dim) + .map(|d| Fgn::new(h[d], steps, Some(t1 - t0), Unseeded).on::()) .collect(); for p in 0..n_paths { diff --git a/stochastic-rs-stochastic/src/traits/process.rs b/stochastic-rs-stochastic/src/traits/process.rs index 8d147b2c..1b6c0336 100644 --- a/stochastic-rs-stochastic/src/traits/process.rs +++ b/stochastic-rs-stochastic/src/traits/process.rs @@ -1,19 +1,5 @@ //! `ProcessExt` and dimensional output markers. -#[cfg(any( - feature = "gpu", - feature = "cuda-native", - feature = "accelerate", - feature = "metal" -))] -use anyhow::Result; -#[cfg(any( - feature = "gpu", - feature = "cuda-native", - feature = "accelerate", - feature = "metal" -))] -use either::Either; use ndarray::Array1; use ndarray::Array2; use ndarray::parallel::prelude::*; @@ -36,23 +22,14 @@ use stochastic_rs_distributions::traits::FloatExt; /// LMM/BGM (see its module doc); it is a parallel array of independent /// Euler-stepped multiplicative martingales. /// -/// ## GPU coverage +/// ## Backend selection /// -/// The `sample_gpu` / `sample_cuda_native` / `sample_metal` / `sample_accelerate` -/// methods are opt-in overrides; default impls return `Err`. Currently the -/// only processes with GPU implementations are [`Fgn`](crate::noise::fgn::Fgn) -/// and [`Fbm`](crate::process::fbm::Fbm), which both rely on the FFT-based -/// circulant-embedding path under `feature = "gpu"`. -/// -/// **Roadmap (P3/18):** GPU implementations for Heston, RoughBergomi, Cir2F, -/// and Hjm remain TODO. These processes use *standard* Brownian noise, -/// not fractional, so the Fgn GPU path is not directly reusable — they -/// need bespoke kernels for variance updates (e.g. Andersen QE for Cir-type -/// dynamics) and correlated noise generation. ([`Bgm`](crate::interest::bgm::Bgm) -/// is excluded from the GPU roadmap on purpose: it is a parallel array of -/// uncoupled Euler-stepped multiplicative martingales, not a market model -/// — see its module doc — so a GPU sampler would just be `xn` independent -/// `Gbm`-style streams and offers no algorithmic interest.) +/// Re-type a process to a compile-time sampling backend with the turbofish +/// `process.on::()` where `B: `[`Backend`](crate::device::Backend) (e.g. +/// `process.on::()`); the backend marker propagates to the +/// process's noise source with no runtime branch. Only the fractional family +/// (built on [`Fgn`](crate::noise::fgn::Fgn)) exposes GPU backends today, and a +/// GPU marker only exists when its feature is compiled. pub trait ProcessExt: Send + Sync { type Output: Send; @@ -61,26 +38,6 @@ pub trait ProcessExt: Send + Sync { fn sample_par(&self, m: usize) -> Vec { (0..m).into_par_iter().map(|_| self.sample()).collect() } - - #[cfg(feature = "gpu")] - fn sample_gpu(&self, _m: usize) -> Result, Array2>> { - anyhow::bail!("CubeCL GPU sampling is not supported for this process") - } - - #[cfg(feature = "cuda-native")] - fn sample_cuda_native(&self, _m: usize) -> Result, Array2>> { - anyhow::bail!("cudarc native CUDA sampling is not supported for this process") - } - - #[cfg(feature = "accelerate")] - fn sample_accelerate(&self, _m: usize) -> Result, Array2>> { - anyhow::bail!("Accelerate/vDSP sampling is not supported for this process") - } - - #[cfg(feature = "metal")] - fn sample_metal(&self, _m: usize) -> Result, Array2>> { - anyhow::bail!("Metal GPU sampling is not supported for this process") - } } /// Marker for processes whose [`ProcessExt::sample`] returns a single diff --git a/tests/device_propagation.rs b/tests/device_propagation.rs new file mode 100644 index 00000000..26b55042 --- /dev/null +++ b/tests/device_propagation.rs @@ -0,0 +1,52 @@ +use stochastic_rs::prelude::*; +use stochastic_rs::simd_rng::Deterministic; +use stochastic_rs::stochastic::device::Cpu; +use stochastic_rs::stochastic::diffusion::fou::Fou; +use stochastic_rs::stochastic::process::fbm::Fbm; + +#[test] +fn fou_on_cpu_matches_plain_sample_with_same_seed() { + let plain = Fou::new( + 0.7_f64, + 1.0, + 0.0, + 0.2, + 64, + Some(0.0), + Some(1.0), + Deterministic::new(7), + ) + .sample(); + let on_cpu = Fou::new( + 0.7_f64, + 1.0, + 0.0, + 0.2, + 64, + Some(0.0), + Some(1.0), + Deterministic::new(7), + ) + .on::() + .sample(); + assert_eq!(plain, on_cpu); +} + +#[test] +fn fbm_on_cpu_matches_plain_sample_with_same_seed() { + let plain = Fbm::new(0.7_f64, 64, Some(1.0), Deterministic::new(11)).sample(); + let on_cpu = Fbm::new(0.7_f64, 64, Some(1.0), Deterministic::new(11)) + .on::() + .sample(); + assert_eq!(plain, on_cpu); +} + +#[test] +fn fbm_sample_par_returns_requested_path_count() { + let m = 8; + let paths = Fbm::new(0.7_f64, 64, Some(1.0), Deterministic::new(3)) + .on::() + .sample_par(m); + assert_eq!(paths.len(), m); + assert!(paths.iter().all(|p| p.len() == 64)); +} diff --git a/tests/fgn_accelerate_validation.rs b/tests/fgn_accelerate_validation.rs index cab1b48d..d56528b0 100644 --- a/tests/fgn_accelerate_validation.rs +++ b/tests/fgn_accelerate_validation.rs @@ -2,22 +2,18 @@ #[cfg(feature = "accelerate")] mod accel_validation { - use either::Either; + use stochastic_rs::simd_rng::Unseeded; + use stochastic_rs::stochastic::device::Accelerate; use stochastic_rs::stochastic::noise::fgn::Fgn; use stochastic_rs::traits::ProcessExt; fn accel_paths(h: f32, n: usize, m: usize) -> Vec> { - let fgn = Fgn::::new(h, n, Some(1.0)); - match fgn - .sample_accelerate(m) - .expect("accelerate sampling failed") - { - Either::Left(p) => vec![p.iter().map(|&x| x as f64).collect()], - Either::Right(ps) => ps - .outer_iter() - .map(|r| r.iter().map(|&x| x as f64).collect()) - .collect(), - } + let fgn = Fgn::::new(h, n, Some(1.0), Unseeded).on::(); + fgn + .sample_par(m) + .into_iter() + .map(|p| p.iter().map(|&x| x as f64).collect()) + .collect() } fn lag_cov(paths: &[Vec], mean: f64, lag: usize) -> f64 { diff --git a/tests/fgn_all_backends_visual.rs b/tests/fgn_all_backends_visual.rs index 9c1da9ad..1b96a422 100644 --- a/tests/fgn_all_backends_visual.rs +++ b/tests/fgn_all_backends_visual.rs @@ -3,62 +3,57 @@ #[cfg(all(feature = "gpu-wgpu", feature = "metal", feature = "accelerate"))] mod all_backends { - use either::Either; use ndarray::Array1; + use stochastic_rs::simd_rng::Unseeded; use stochastic_rs::stats::fractal_dim::FractalDimEstimator; use stochastic_rs::stats::fractal_dim::Higuchi; + use stochastic_rs::stochastic::device::Accelerate; + use stochastic_rs::stochastic::device::CubeCl; + use stochastic_rs::stochastic::device::MetalNative; use stochastic_rs::stochastic::noise::fgn::Fgn; use stochastic_rs::stochastic::process::fbm::Fbm; use stochastic_rs::traits::ProcessExt; use stochastic_rs::visualization::GridPlotter; fn cpu_fgn(h: f64, n: usize, m: usize) -> Vec> { - let fgn = Fgn::::new(h, n, Some(1.0)); + let fgn = Fgn::::new(h, n, Some(1.0), Unseeded); fgn.sample_par(m).into_iter().map(|p| p.to_vec()).collect() } fn gpu_fgn(h: f32, n: usize, m: usize) -> Vec> { - let fgn = Fgn::::new(h, n, Some(1.0)); - match fgn.sample_gpu(m).unwrap() { - Either::Left(p) => vec![p.iter().map(|&x| x as f64).collect()], - Either::Right(ps) => ps - .outer_iter() - .map(|r| r.iter().map(|&x| x as f64).collect()) - .collect(), - } + let fgn = Fgn::::new(h, n, Some(1.0), Unseeded).on::(); + fgn + .sample_par(m) + .into_iter() + .map(|p| p.iter().map(|&x| x as f64).collect()) + .collect() } fn metal_fgn(h: f32, n: usize, m: usize) -> Vec> { - let fgn = Fgn::::new(h, n, Some(1.0)); - match fgn.sample_metal(m).unwrap() { - Either::Left(p) => vec![p.iter().map(|&x| x as f64).collect()], - Either::Right(ps) => ps - .outer_iter() - .map(|r| r.iter().map(|&x| x as f64).collect()) - .collect(), - } + let fgn = Fgn::::new(h, n, Some(1.0), Unseeded).on::(); + fgn + .sample_par(m) + .into_iter() + .map(|p| p.iter().map(|&x| x as f64).collect()) + .collect() } fn accel_fgn(h: f32, n: usize, m: usize) -> Vec> { - let fgn = Fgn::::new(h, n, Some(1.0)); - match fgn.sample_accelerate(m).unwrap() { - Either::Left(p) => vec![p.iter().map(|&x| x as f64).collect()], - Either::Right(ps) => ps - .outer_iter() - .map(|r| r.iter().map(|&x| x as f64).collect()) - .collect(), - } + let fgn = Fgn::::new(h, n, Some(1.0), Unseeded).on::(); + fgn + .sample_par(m) + .into_iter() + .map(|p| p.iter().map(|&x| x as f64).collect()) + .collect() } fn metal_fbm(h: f32, n: usize, m: usize) -> Vec> { - let fbm = Fbm::::new(h, n, Some(1.0)); - match fbm.sample_metal(m).unwrap() { - Either::Left(p) => vec![p.iter().map(|&x| x as f64).collect()], - Either::Right(ps) => ps - .outer_iter() - .map(|r| r.iter().map(|&x| x as f64).collect()) - .collect(), - } + let fbm = Fbm::::new(h, n, Some(1.0), Unseeded).on::(); + fbm + .sample_par(m) + .into_iter() + .map(|p| p.iter().map(|&x| x as f64).collect()) + .collect() } fn empirical_acov(paths: &[Vec], max_lag: usize) -> Vec { diff --git a/tests/fgn_cuda_oxide_validation.rs b/tests/fgn_cuda_oxide_validation.rs new file mode 100644 index 00000000..0d1d72d9 --- /dev/null +++ b/tests/fgn_cuda_oxide_validation.rs @@ -0,0 +1,64 @@ +//! Validates the experimental cuda-oxide FGN backend against covariance +//! structure. Requires building/running through `cargo oxide`. + +#[cfg(feature = "cuda-oxide-experimental")] +mod cuda_oxide_validation { + use stochastic_rs::simd_rng::Unseeded; + use stochastic_rs::stochastic::noise::fgn::Fgn; + + fn lag_covariance(paths: &[Vec], mean: f64, lag: usize) -> f64 { + let mut s = 0.0; + let mut c = 0usize; + for p in paths { + for i in 0..(p.len() - lag) { + s += (p[i] - mean) * (p[i + lag] - mean); + c += 1; + } + } + s / c as f64 + } + + fn cuda_oxide_paths(h: f32, n: usize, m: usize) -> Vec> { + let fgn = Fgn::::new(h, n, Some(1.0), Unseeded); + fgn + .sample_cuda_oxide_with_module(m, "fgn_cuda_oxide_validation") + .expect("cuda-oxide sampling failed") + .outer_iter() + .map(|row| row.iter().map(|&x| x as f64).collect()) + .collect() + } + + #[test] + fn cuda_oxide_fgn_variance_and_lag_sign() { + let h = 0.72_f32; + let n = 512_usize; + let m = 2048_usize; + let paths = cuda_oxide_paths(h, n, m); + + let values: Vec = paths.iter().flatten().copied().collect(); + let mean = values.iter().sum::() / values.len() as f64; + let var = values.iter().map(|x| (x - mean).powi(2)).sum::() / values.len() as f64; + let var_theory = (1.0 / n as f64).powf(2.0 * h as f64); + let cov1 = lag_covariance(&paths, mean, 1); + + assert!(mean.abs() < 0.01, "mean too far from zero: {mean}"); + assert!( + ((var / var_theory) - 1.0).abs() < 0.15, + "variance ratio: {}", + var / var_theory + ); + assert!( + cov1 > 0.0, + "H>0.5 should have positive lag-1 covariance: {cov1}" + ); + } + + #[test] + fn cuda_oxide_non_power_of_two_n_shape() { + let fgn = Fgn::::new(0.7, 3000, Some(1.0), Unseeded); + let batch = fgn + .sample_cuda_oxide_with_module(8, "fgn_cuda_oxide_validation") + .expect("cuda-oxide batch"); + assert_eq!(batch.shape(), &[8, 3000]); + } +} diff --git a/tests/fgn_gpu_validation.rs b/tests/fgn_gpu_validation.rs index d5e1d246..f93f4bd6 100644 --- a/tests/fgn_gpu_validation.rs +++ b/tests/fgn_gpu_validation.rs @@ -7,9 +7,10 @@ //! If the GPU FFT is correct, empirical covariance from GPU samples must //! match this formula to within Monte Carlo noise. -#[cfg(feature = "gpu-wgpu")] +#[cfg(any(feature = "gpu-cuda", feature = "gpu-wgpu"))] mod gpu_fft_validation { - use either::Either; + use stochastic_rs::simd_rng::Unseeded; + use stochastic_rs::stochastic::device::CubeCl; use stochastic_rs::stochastic::noise::fgn::Fgn; use stochastic_rs::traits::ProcessExt; @@ -36,18 +37,16 @@ mod gpu_fft_validation { } fn sample_gpu_paths(h: f32, n: usize, t: f32, m: usize) -> Vec> { - let fgn = Fgn::::new(h, n, Some(t)); - match fgn.sample_gpu(m).expect("GPU sampling failed") { - Either::Left(path) => vec![path.iter().map(|&x| x as f64).collect()], - Either::Right(paths) => paths - .outer_iter() - .map(|row| row.iter().map(|&x| x as f64).collect()) - .collect(), - } + let fgn = Fgn::::new(h, n, Some(t), Unseeded).on::(); + fgn + .sample_par(m) + .into_iter() + .map(|path| path.iter().map(|&x| x as f64).collect()) + .collect() } fn sample_cpu_paths(h: f64, n: usize, t: f64, m: usize) -> Vec> { - let fgn = Fgn::::new(h, n, Some(t)); + let fgn = Fgn::::new(h, n, Some(t), Unseeded); fgn .sample_par(m) .into_iter() diff --git a/tests/fgn_gpu_visual.rs b/tests/fgn_gpu_visual.rs index df513d31..22a60d6a 100644 --- a/tests/fgn_gpu_visual.rs +++ b/tests/fgn_gpu_visual.rs @@ -4,43 +4,40 @@ //! trajectories for multiple Hurst parameters. Also estimates H back //! from the generated fBM paths using fractal dimension. -#[cfg(feature = "gpu-wgpu")] +#[cfg(any(feature = "gpu-cuda", feature = "gpu-wgpu"))] mod gpu_visual { - use either::Either; use ndarray::Array1; + use stochastic_rs::simd_rng::Unseeded; use stochastic_rs::stats::fractal_dim::FractalDimEstimator; use stochastic_rs::stats::fractal_dim::Higuchi; use stochastic_rs::stats::fractal_dim::Variogram; + use stochastic_rs::stochastic::device::CubeCl; use stochastic_rs::stochastic::noise::fgn::Fgn; use stochastic_rs::stochastic::process::fbm::Fbm; use stochastic_rs::traits::ProcessExt; use stochastic_rs::visualization::GridPlotter; fn gpu_fgn_paths(h: f32, n: usize, m: usize) -> Vec> { - let fgn = Fgn::::new(h, n, Some(1.0)); - match fgn.sample_gpu(m).expect("GPU Fgn sampling failed") { - Either::Left(p) => vec![p.iter().map(|&x| x as f64).collect()], - Either::Right(ps) => ps - .outer_iter() - .map(|r| r.iter().map(|&x| x as f64).collect()) - .collect(), - } + let fgn = Fgn::::new(h, n, Some(1.0), Unseeded).on::(); + fgn + .sample_par(m) + .into_iter() + .map(|p| p.iter().map(|&x| x as f64).collect()) + .collect() } fn cpu_fgn_paths(h: f64, n: usize, m: usize) -> Vec> { - let fgn = Fgn::::new(h, n, Some(1.0)); + let fgn = Fgn::::new(h, n, Some(1.0), Unseeded); fgn.sample_par(m).into_iter().map(|p| p.to_vec()).collect() } fn gpu_fbm_paths(h: f32, n: usize, m: usize) -> Vec> { - let fbm = Fbm::::new(h, n, Some(1.0)); - match fbm.sample_gpu(m).expect("GPU Fbm sampling failed") { - Either::Left(p) => vec![p.iter().map(|&x| x as f64).collect()], - Either::Right(ps) => ps - .outer_iter() - .map(|r| r.iter().map(|&x| x as f64).collect()) - .collect(), - } + let fbm = Fbm::::new(h, n, Some(1.0), Unseeded).on::(); + fbm + .sample_par(m) + .into_iter() + .map(|p| p.iter().map(|&x| x as f64).collect()) + .collect() } fn empirical_autocovariance(paths: &[Vec], max_lag: usize) -> Vec { diff --git a/tests/fgn_metal_validation.rs b/tests/fgn_metal_validation.rs index 56ddc9ae..c5c41842 100644 --- a/tests/fgn_metal_validation.rs +++ b/tests/fgn_metal_validation.rs @@ -1,19 +1,17 @@ #[cfg(feature = "metal")] mod metal_validation { - use either::Either; use stochastic_rs::simd_rng::Unseeded; + use stochastic_rs::stochastic::device::MetalNative; use stochastic_rs::stochastic::noise::fgn::Fgn; use stochastic_rs::traits::ProcessExt; fn metal_paths(h: f32, n: usize, m: usize) -> Vec> { - let fgn = Fgn::::new(h, n, Some(1.0), Unseeded); - match fgn.sample_metal(m).expect("metal sampling failed") { - Either::Left(p) => vec![p.iter().map(|&x| x as f64).collect()], - Either::Right(ps) => ps - .outer_iter() - .map(|r| r.iter().map(|&x| x as f64).collect()) - .collect(), - } + let fgn = Fgn::::new(h, n, Some(1.0), Unseeded).on::(); + fgn + .sample_par(m) + .into_iter() + .map(|p| p.iter().map(|&x| x as f64).collect()) + .collect() } fn lag_cov(paths: &[Vec], mean: f64, lag: usize) -> f64 {