diffusionx 0.11.2

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

#[cfg(any(feature = "cuda", feature = "metal"))]
use std::{env, path::PathBuf, process::Command};

#[cfg(not(any(feature = "cuda", feature = "metal")))]
fn main() {}

#[cfg(all(feature = "cuda", feature = "metal"))]
fn main() {
    panic!("Cannot enable both CUDA and Metal features");
}

#[cfg(feature = "cuda")]
fn main() {
    println!("cargo:rerun-if-changed=build.rs");

    let kernels = ["bm", "random", "levy", "ou"];

    let out_dir = PathBuf::from(env::var("OUT_DIR").expect("OUT_DIR not set"));

    for kernel in kernels {
        let kernel_file = format!("kernels/cuda-kernel/{kernel}.cu");
        println!("cargo:rerun-if-changed={kernel_file}");

        let ptx_path = out_dir.join(format!("{kernel}.ptx"));

        let status = Command::new("nvcc")
            .args([
                "-ptx",
                &kernel_file,
                "-o",
                ptx_path.to_str().expect("Invalid PTX path"),
            ])
            .status()
            .expect("Failed to execute nvcc. Is CUDA installed and nvcc in PATH?");

        if !status.success() {
            panic!("nvcc failed with status: {status}");
        }

        let upper_kernel = kernel.to_uppercase();

        println!(
            "cargo:rustc-env={upper_kernel}_KERNEL_PTX={}",
            ptx_path.display()
        );
    }
}

#[cfg(feature = "metal")]
fn main() {
    println!("cargo:rerun-if-changed=build.rs");

    let kernels = ["bm", "levy", "ou", "random"];
    let kernel_dir = "kernels/metal-kernel";

    let out_dir = PathBuf::from(env::var("OUT_DIR").expect("OUT_DIR not set"));

    // Read utils.metal content for inclusion
    let utils_path = format!("{kernel_dir}/utils.metal");
    println!("cargo:rerun-if-changed={utils_path}");

    for kernel in kernels {
        let kernel_file = format!("{kernel_dir}/{kernel}.metal");
        println!("cargo:rerun-if-changed={kernel_file}");

        let metallib_path = out_dir.join(format!("{kernel}.metallib"));
        let air_path = out_dir.join(format!("{kernel}.air"));

        // Compile .metal to .air (Apple Intermediate Representation)
        let status = Command::new("xcrun")
            .args([
                "-sdk",
                "macosx",
                "metal",
                "-c",
                &kernel_file,
                "-I",
                kernel_dir, // Include directory for utils.metal
                "-o",
                air_path.to_str().expect("Invalid AIR path"),
            ])
            .status()
            .expect("Failed to execute xcrun metal. Is Xcode installed?");

        if !status.success() {
            panic!("Metal compiler failed with status: {status}");
        }

        // Link .air to .metallib
        let status = Command::new("xcrun")
            .args([
                "-sdk",
                "macosx",
                "metallib",
                air_path.to_str().expect("Invalid AIR path"),
                "-o",
                metallib_path.to_str().expect("Invalid metallib path"),
            ])
            .status()
            .expect("Failed to execute xcrun metallib");

        if !status.success() {
            panic!("Metal linker failed with status: {status}");
        }

        let upper_kernel = kernel.to_uppercase();

        println!(
            "cargo:rustc-env={upper_kernel}_KERNEL_METALLIB={}",
            metallib_path.display()
        );
    }
}