llama_cpp_2/
llama_backend.rs

1//! Representation of an initialized llama backend
2
3use crate::LlamaCppError;
4use llama_cpp_sys_2::ggml_log_level;
5use std::sync::atomic::AtomicBool;
6use std::sync::atomic::Ordering::SeqCst;
7
8/// Representation of an initialized llama backend
9/// This is required as a parameter for most llama functions as the backend must be initialized
10/// before any llama functions are called. This type is proof of initialization.
11#[derive(Eq, PartialEq, Debug)]
12pub struct LlamaBackend {}
13
14static LLAMA_BACKEND_INITIALIZED: AtomicBool = AtomicBool::new(false);
15
16impl LlamaBackend {
17    /// Mark the llama backend as initialized
18    fn mark_init() -> crate::Result<()> {
19        match LLAMA_BACKEND_INITIALIZED.compare_exchange(false, true, SeqCst, SeqCst) {
20            Ok(_) => Ok(()),
21            Err(_) => Err(LlamaCppError::BackendAlreadyInitialized),
22        }
23    }
24
25    /// Initialize the llama backend (without numa).
26    ///
27    /// # Examples
28    ///
29    /// ```
30    ///# use llama_cpp_2::llama_backend::LlamaBackend;
31    ///# use llama_cpp_2::LlamaCppError;
32    ///# use std::error::Error;
33    ///
34    ///# fn main() -> Result<(), Box<dyn Error>> {
35    ///
36    ///
37    /// let backend = LlamaBackend::init()?;
38    /// // the llama backend can only be initialized once
39    /// assert_eq!(Err(LlamaCppError::BackendAlreadyInitialized), LlamaBackend::init());
40    ///
41    ///# Ok(())
42    ///# }
43    /// ```
44    #[tracing::instrument(skip_all)]
45    pub fn init() -> crate::Result<LlamaBackend> {
46        Self::mark_init()?;
47        unsafe { llama_cpp_sys_2::llama_backend_init() }
48        Ok(LlamaBackend {})
49    }
50
51    /// Initialize the llama backend (with numa).
52    /// ```
53    ///# use llama_cpp_2::llama_backend::LlamaBackend;
54    ///# use std::error::Error;
55    ///# use llama_cpp_2::llama_backend::NumaStrategy;
56    ///
57    ///# fn main() -> Result<(), Box<dyn Error>> {
58    ///
59    /// let llama_backend = LlamaBackend::init_numa(NumaStrategy::MIRROR)?;
60    ///
61    ///# Ok(())
62    ///# }
63    /// ```
64    #[tracing::instrument(skip_all)]
65    pub fn init_numa(strategy: NumaStrategy) -> crate::Result<LlamaBackend> {
66        // Initialize backends first, then NUMA.
67        let this = Self::init()?;
68        unsafe {
69            llama_cpp_sys_2::llama_numa_init(llama_cpp_sys_2::ggml_numa_strategy::from(strategy));
70        }
71        Ok(this)
72    }
73
74    /// Was the code built for a GPU backend & is a supported one available.
75    pub fn supports_gpu_offload(&self) -> bool {
76        unsafe { llama_cpp_sys_2::llama_supports_gpu_offload() }
77    }
78
79    /// Does this platform support loading the model via mmap.
80    pub fn supports_mmap(&self) -> bool {
81        unsafe { llama_cpp_sys_2::llama_supports_mmap() }
82    }
83
84    /// Does this platform support locking the model in RAM.
85    pub fn supports_mlock(&self) -> bool {
86        unsafe { llama_cpp_sys_2::llama_supports_mlock() }
87    }
88
89    /// Change the output of llama.cpp's logging to be voided instead of pushed to `stderr`.
90    pub fn void_logs(&mut self) {
91        unsafe extern "C" fn void_log(
92            _level: ggml_log_level,
93            _text: *const ::std::os::raw::c_char,
94            _user_data: *mut ::std::os::raw::c_void,
95        ) {
96        }
97
98        unsafe {
99            llama_cpp_sys_2::llama_log_set(Some(void_log), std::ptr::null_mut());
100        }
101    }
102}
103
104/// A rusty wrapper around `numa_strategy`.
105#[derive(Debug, Eq, PartialEq, Copy, Clone)]
106pub enum NumaStrategy {
107    /// The numa strategy is disabled.
108    DISABLED,
109    /// help wanted: what does this do?
110    DISTRIBUTE,
111    /// help wanted: what does this do?
112    ISOLATE,
113    /// help wanted: what does this do?
114    NUMACTL,
115    /// help wanted: what does this do?
116    MIRROR,
117    /// help wanted: what does this do?
118    COUNT,
119}
120
121/// An invalid numa strategy was provided.
122#[derive(Debug, Eq, PartialEq, Copy, Clone)]
123pub struct InvalidNumaStrategy(
124    /// The invalid numa strategy that was provided.
125    pub llama_cpp_sys_2::ggml_numa_strategy,
126);
127
128impl TryFrom<llama_cpp_sys_2::ggml_numa_strategy> for NumaStrategy {
129    type Error = InvalidNumaStrategy;
130
131    fn try_from(value: llama_cpp_sys_2::ggml_numa_strategy) -> Result<Self, Self::Error> {
132        match value {
133            llama_cpp_sys_2::GGML_NUMA_STRATEGY_DISABLED => Ok(Self::DISABLED),
134            llama_cpp_sys_2::GGML_NUMA_STRATEGY_DISTRIBUTE => Ok(Self::DISTRIBUTE),
135            llama_cpp_sys_2::GGML_NUMA_STRATEGY_ISOLATE => Ok(Self::ISOLATE),
136            llama_cpp_sys_2::GGML_NUMA_STRATEGY_NUMACTL => Ok(Self::NUMACTL),
137            llama_cpp_sys_2::GGML_NUMA_STRATEGY_MIRROR => Ok(Self::MIRROR),
138            llama_cpp_sys_2::GGML_NUMA_STRATEGY_COUNT => Ok(Self::COUNT),
139            value => Err(InvalidNumaStrategy(value)),
140        }
141    }
142}
143
144impl From<NumaStrategy> for llama_cpp_sys_2::ggml_numa_strategy {
145    fn from(value: NumaStrategy) -> Self {
146        match value {
147            NumaStrategy::DISABLED => llama_cpp_sys_2::GGML_NUMA_STRATEGY_DISABLED,
148            NumaStrategy::DISTRIBUTE => llama_cpp_sys_2::GGML_NUMA_STRATEGY_DISTRIBUTE,
149            NumaStrategy::ISOLATE => llama_cpp_sys_2::GGML_NUMA_STRATEGY_ISOLATE,
150            NumaStrategy::NUMACTL => llama_cpp_sys_2::GGML_NUMA_STRATEGY_NUMACTL,
151            NumaStrategy::MIRROR => llama_cpp_sys_2::GGML_NUMA_STRATEGY_MIRROR,
152            NumaStrategy::COUNT => llama_cpp_sys_2::GGML_NUMA_STRATEGY_COUNT,
153        }
154    }
155}
156
157/// Drops the llama backend.
158/// ```
159///
160///# use llama_cpp_2::llama_backend::LlamaBackend;
161///# use std::error::Error;
162///
163///# fn main() -> Result<(), Box<dyn Error>> {
164/// let backend = LlamaBackend::init()?;
165/// drop(backend);
166/// // can be initialized again after being dropped
167/// let backend = LlamaBackend::init()?;
168///# Ok(())
169///# }
170///
171/// ```
172impl Drop for LlamaBackend {
173    fn drop(&mut self) {
174        match LLAMA_BACKEND_INITIALIZED.compare_exchange(true, false, SeqCst, SeqCst) {
175            Ok(_) => {}
176            Err(_) => {
177                unreachable!("This should not be reachable as the only ways to obtain a llama backend involve marking the backend as initialized.")
178            }
179        }
180        unsafe { llama_cpp_sys_2::llama_backend_free() }
181    }
182}
183
184/// Load GGML backend modules from the given directory.
185///
186/// Call this before [`LlamaBackend::init`] to enable runtime hardware selection
187/// (Vulkan, CPU-AVX512, CPU-AVX2, etc.) when built with the `dynamic-backends` feature.
188#[cfg(feature = "dynamic-backends")]
189pub fn load_backends_from_path(path: &std::path::Path) {
190    let s = std::ffi::CString::new(path.to_str().expect("path must be valid UTF-8"))
191        .expect("path must not contain null bytes");
192    unsafe { llama_cpp_sys_2::ggml_backend_load_all_from_path(s.as_ptr()) }
193}
194
195/// Load GGML backend modules from the compile-time default directory ([`BACKENDS_DIR`]).
196///
197/// This is a no-op when `BACKENDS_DIR` is `None` (static builds or development builds
198/// that have not set `GGML_BACKENDS_DIR`).
199#[cfg(feature = "dynamic-backends")]
200pub fn load_backends() {
201    if let Some(dir) = llama_cpp_sys_2::BACKENDS_DIR {
202        load_backends_from_path(std::path::Path::new(dir));
203    }
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209
210    #[test]
211    fn numa_from_and_to() {
212        let numas = [
213            NumaStrategy::DISABLED,
214            NumaStrategy::DISTRIBUTE,
215            NumaStrategy::ISOLATE,
216            NumaStrategy::NUMACTL,
217            NumaStrategy::MIRROR,
218            NumaStrategy::COUNT,
219        ];
220
221        for numa in &numas {
222            let from = llama_cpp_sys_2::ggml_numa_strategy::from(*numa);
223            let to = NumaStrategy::try_from(from).expect("Failed to convert from and to");
224            assert_eq!(*numa, to);
225        }
226    }
227
228    #[test]
229    fn check_invalid_numa() {
230        let invalid = 800;
231        let invalid = NumaStrategy::try_from(invalid);
232        assert_eq!(invalid, Err(InvalidNumaStrategy(invalid.unwrap_err().0)));
233    }
234}
llama_cpp_2/llama_backend.rs

llama_cpp_2/
llama_backend.rs
