llama_cpp_bindings/
llama_backend.rs

1use std::sync::atomic::AtomicBool;
2use std::sync::atomic::Ordering::SeqCst;
3
4use llama_cpp_bindings_sys::ggml_log_level;
5
6use crate::LlamaCppError;
7use crate::llama_backend_numa_strategy::NumaStrategy;
8
9#[derive(Eq, PartialEq, Debug)]
10pub struct LlamaBackend {}
11
12static LLAMA_BACKEND_INITIALIZED: AtomicBool = AtomicBool::new(false);
13
14impl LlamaBackend {
15    fn mark_init() -> crate::Result<()> {
16        match LLAMA_BACKEND_INITIALIZED.compare_exchange(false, true, SeqCst, SeqCst) {
17            Ok(_was_uninitialized) => Ok(()),
18            Err(_was_already_initialized) => Err(LlamaCppError::BackendAlreadyInitialized),
19        }
20    }
21
22    /// # Errors
23    /// Returns an error if the backend was already initialized.
24    pub fn init() -> crate::Result<Self> {
25        Self::mark_init()?;
26        unsafe { llama_cpp_bindings_sys::llama_backend_init() }
27        Ok(Self {})
28    }
29
30    /// # Errors
31    /// Returns an error if the backend was already initialized.
32    pub fn init_numa(strategy: NumaStrategy) -> crate::Result<Self> {
33        Self::mark_init()?;
34        unsafe {
35            llama_cpp_bindings_sys::llama_numa_init(
36                llama_cpp_bindings_sys::ggml_numa_strategy::from(strategy),
37            );
38        }
39        Ok(Self {})
40    }
41
42    #[must_use]
43    pub fn supports_gpu_offload(&self) -> bool {
44        unsafe { llama_cpp_bindings_sys::llama_supports_gpu_offload() }
45    }
46
47    #[must_use]
48    pub fn supports_mmap(&self) -> bool {
49        unsafe { llama_cpp_bindings_sys::llama_supports_mmap() }
50    }
51
52    #[must_use]
53    pub fn supports_mlock(&self) -> bool {
54        unsafe { llama_cpp_bindings_sys::llama_supports_mlock() }
55    }
56
57    pub fn void_logs(&mut self) {
58        unsafe {
59            llama_cpp_bindings_sys::llama_log_set(Some(void_log), std::ptr::null_mut());
60        }
61    }
62}
63
64const unsafe extern "C" fn void_log(
65    _level: ggml_log_level,
66    _text: *const ::std::os::raw::c_char,
67    _user_data: *mut ::std::os::raw::c_void,
68) {
69}
70
71impl Drop for LlamaBackend {
72    fn drop(&mut self) {
73        LLAMA_BACKEND_INITIALIZED.store(false, SeqCst);
74        unsafe { llama_cpp_bindings_sys::llama_backend_free() }
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use serial_test::serial;
81
82    use super::LlamaBackend;
83    use crate::LlamaCppError;
84
85    #[test]
86    fn void_log_callback_does_not_panic() {
87        unsafe {
88            super::void_log(
89                llama_cpp_bindings_sys::GGML_LOG_LEVEL_INFO,
90                c"test".as_ptr(),
91                std::ptr::null_mut(),
92            );
93        }
94    }
95
96    #[test]
97    #[serial]
98    fn init_succeeds() {
99        let backend = LlamaBackend::init();
100        assert!(backend.is_ok());
101    }
102
103    #[test]
104    #[serial]
105    fn double_init_returns_error() {
106        let _backend = LlamaBackend::init().unwrap();
107        let second_err = LlamaBackend::init().unwrap_err();
108
109        assert_eq!(
110            std::mem::discriminant(&second_err),
111            std::mem::discriminant(&LlamaCppError::BackendAlreadyInitialized),
112            "expected BackendAlreadyInitialized, got {second_err:?}"
113        );
114    }
115
116    #[test]
117    #[serial]
118    fn init_numa_returns_error_when_backend_already_initialized() {
119        use crate::llama_backend_numa_strategy::NumaStrategy;
120
121        let _backend = LlamaBackend::init().unwrap();
122        let second_err = LlamaBackend::init_numa(NumaStrategy::Disabled).unwrap_err();
123
124        assert_eq!(
125            std::mem::discriminant(&second_err),
126            std::mem::discriminant(&LlamaCppError::BackendAlreadyInitialized),
127            "expected BackendAlreadyInitialized, got {second_err:?}"
128        );
129    }
130
131    #[test]
132    #[serial]
133    fn feature_queries_return_bools() {
134        let backend = LlamaBackend::init().unwrap();
135        let _gpu = backend.supports_gpu_offload();
136        let _mmap = backend.supports_mmap();
137        let _mlock = backend.supports_mlock();
138    }
139
140    #[test]
141    #[serial]
142    fn drop_and_reinit_works() {
143        let backend = LlamaBackend::init().unwrap();
144        drop(backend);
145        let backend = LlamaBackend::init();
146        assert!(backend.is_ok());
147    }
148
149    #[test]
150    #[serial]
151    fn init_numa_succeeds() {
152        use crate::llama_backend_numa_strategy::NumaStrategy;
153
154        let backend = LlamaBackend::init_numa(NumaStrategy::Disabled);
155        assert!(backend.is_ok());
156    }
157}
llama_cpp_bindings/llama_backend.rs

llama_cpp_bindings/
llama_backend.rs
