multiscreen_rs/
device.rs

1//! Device abstraction for CPU and CUDA.
2
3#[cfg(not(feature = "cuda"))]
4use crate::error::Error;
5use crate::error::Result;
6use crate::runtime::Device;
7
8/// Returns the default CPU device.
9///
10/// ```rust,no_run
11/// use multiscreen_rs::prelude::*;
12/// fn main() -> multiscreen_rs::Result<()> {
13///     let device = cpu()?;
14///     Ok(())
15/// }
16/// ```
17pub fn cpu() -> Result<Device> {
18    Ok(Device::default())
19}
20
21/// Returns a CUDA device for the given GPU index.
22///
23/// Only available with the `cuda` feature. Returns a clear error if CUDA
24/// is not available.
25///
26/// ```toml
27/// [dependencies]
28/// multiscreen-rs = { version = "0.1", features = ["cuda"] }
29/// ```
30#[cfg(feature = "cuda")]
31pub fn cuda(index: usize) -> Result<crate::runtime::CudaDevice> {
32    // For now, Burn's Cuda device doesn't support selecting GPU index via
33    // the simple API, so we just return the default CUDA device.
34    // The index parameter is reserved for future multi-GPU support.
35    let _ = index;
36    Ok(crate::runtime::CudaDevice::default())
37}
38
39/// Returns a CUDA device for the given GPU index.
40///
41/// Without the `cuda` feature, this always returns an error.
42#[cfg(not(feature = "cuda"))]
43pub fn cuda(_index: usize) -> Result<Device> {
44    Err(Error::Config(
45        "CUDA is not available. Enable the 'cuda' feature in Cargo.toml:\n\
46         [dependencies]\n\
47         multiscreen-rs = { version = \"0.1\", features = [\"cuda\"] }"
48            .to_string(),
49    ))
50}
51
52/// Returns the best available device.
53///
54/// Always returns the default Flex (CPU) device. To use CUDA, construct
55/// a [`MultiscreenModel`](crate::MultiscreenModel) with
56/// [`CudaAutodiffBackend`](crate::runtime::CudaAutodiffBackend) directly.
57pub fn auto_device() -> Result<Device> {
58    cpu()
59}
multiscreen_rs/device.rs

multiscreen_rs/
device.rs
