llama_cpp_bindings/
llama_backend.rs

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

llama_cpp_bindings/
llama_backend.rs
