fastnn 0.3.3

A robust deep learning library in Rust with CUDA kernels
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

use std::env;
use std::path::PathBuf;

fn main() {
    println!("cargo:rerun-if-changed=cuda/");

    if cfg!(feature = "cuda") {
        build_cuda();
    } else {
        build_cpu_stubs();
    }
}

fn build_cpu_stubs() {
    cc::Build::new()
        .file("cuda/stubs.c")
        .compile("fastnn_cuda_stubs");
}

fn build_cuda() {
    let cuda_path = env::var("CUDA_PATH")
        .or_else(|_| env::var("CUDA_HOME"))
        .unwrap_or_else(|_| {
            if cfg!(target_os = "windows") {
                "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.0".to_string()
            } else {
                "/usr/local/cuda".to_string()
            }
        });

    let cuda_include = PathBuf::from(&cuda_path).join("include");
    let cuda_lib = if cfg!(target_os = "windows") {
        PathBuf::from(&cuda_path).join("lib/x64")
    } else {
        PathBuf::from(&cuda_path).join("lib64")
    };

    let mut build = cc::Build::new();
    build
        .cuda(true)
        .cudart("shared")
        .flag("-gencode=arch=compute_75,code=sm_75") // Turing
        .flag("-gencode=arch=compute_80,code=sm_80") // Ampere
        .flag("-gencode=arch=compute_86,code=sm_86") // Ampere (RTX 30xx)
        .flag("-gencode=arch=compute_89,code=sm_89") // Ada Lovelace
        .flag("-gencode=arch=compute_90,code=sm_90") // Hopper
        .flag("--use_fast_math")
        .flag("--extended-lambda")
        .flag("-O3")
        .include("cuda/include")
        .include(&cuda_include)
        .file("cuda/kernels.cu");

    // nvcc rejects host compilers newer than it knows. Take FASTNN_NVCC_CCBIN, else g++-15 if installed.
    println!("cargo:rerun-if-env-changed=FASTNN_NVCC_CCBIN");
    if let Ok(ccbin) = env::var("FASTNN_NVCC_CCBIN") {
        build.flag(format!("-ccbin={ccbin}"));
    } else if std::process::Command::new("g++-15")
        .arg("--version")
        .output()
        .is_ok()
    {
        build.flag("-ccbin=g++-15");
    }

    build.compile("fastnn_cuda_kernels");

    println!("cargo:rustc-link-search=native={}", cuda_lib.display());
    println!("cargo:rustc-link-lib=dylib=cudart");
    println!("cargo:rustc-link-lib=dylib=cublas");
    println!("cargo:rustc-link-lib=dylib=curand");
}