fastnn 0.3.1

A small deep learning library in Rust with optional 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

//! GPU layer: device setup, memory ([`buffer`]), raw bindings ([`ffi`]) and safe wrappers ([`kernels`]).
//!
//! Without the `cuda` feature, `cuda/stubs.c` stands in for the kernels and [`device_count`] is 0.

pub mod buffer;
pub mod ffi;
pub mod kernels;

pub use buffer::{empty_cache, CudaBuffer};

use std::sync::OnceLock;

use crate::error::{Error, Result};

/// Sets up the CUDA context on `device`. Only the first call does anything.
pub fn init(device: usize) -> Result<()> {
    static INIT: OnceLock<std::result::Result<(), String>> = OnceLock::new();
    INIT.get_or_init(|| {
        let code = unsafe { ffi::fastnn_cuda_init(device as i32) };
        if code == 0 {
            Ok(())
        } else {
            Err(format!("could not initialise device {device}"))
        }
    })
    .clone()
    .map_err(Error::Cuda)
}

/// Number of visible CUDA devices, 0 in a CPU-only build.
pub fn device_count() -> usize {
    unsafe { ffi::fastnn_cuda_device_count().max(0) as usize }
}

/// Waits for every queued kernel to finish.
pub fn synchronize() -> Result<()> {
    match unsafe { ffi::fastnn_cuda_synchronize() } {
        0 => Ok(()),
        code => Err(Error::Cuda(format!("synchronize failed (status {code})"))),
    }
}

/// `(free, total)` device memory in bytes.
pub fn memory_info() -> (usize, usize) {
    let (mut free, mut total) = (0usize, 0usize);
    unsafe { ffi::fastnn_cuda_get_memory_info(&mut free, &mut total) };
    (free, total)
}