Skip to main content

llama_server/
backend.rs

1//! Pick your preferred compute backends.
2use bitflags::bitflags;
3
4/// A compute backend.
5#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
6pub enum Backend {
7    /// The NVIDIA CUDA backend.
8    Cuda,
9    /// The AMD HIP backend.
10    Hip,
11}
12
13bitflags! {
14    /// A set of compute backends.
15    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
16    pub struct Set: u32 {
17        /// The NVIDIA CUDA backend.
18        const CUDA = 1;
19        /// The AMD HIP backend.
20        const HIP = 1 << 1;
21    }
22}
23
24impl Set {
25    /// Returns the backends in the [`Set`] that are also available in the current
26    /// platform.
27    pub fn available(self) -> impl Iterator<Item = Backend> {
28        let mut backends = Vec::new();
29
30        if cfg!(target_os = "macos") {
31            return backends.into_iter();
32        }
33
34        if self.contains(Self::CUDA) {
35            backends.push(Backend::Cuda);
36        }
37
38        if self.contains(Self::HIP) {
39            backends.push(Backend::Hip);
40        }
41
42        backends.into_iter()
43    }
44
45    /// Returns a new [`Set`] with any unavailable backends filtered out.
46    pub fn normalize(self) -> Self {
47        self.available().fold(Self::empty(), |backends, backend| {
48            backends
49                | match backend {
50                    Backend::Cuda => Self::CUDA,
51                    Backend::Hip => Self::HIP,
52                }
53        })
54    }
55}