Skip to main content

llama_cpp_bindings/
llama_backend_numa_strategy.rs

1use crate::invalid_numa_strategy::InvalidNumaStrategy;
2
3/// NUMA (Non-Uniform Memory Access) thread affinity strategy for llama.cpp.
4#[derive(Debug, Eq, PartialEq, Copy, Clone)]
5pub enum NumaStrategy {
6    /// NUMA-aware scheduling is disabled. Threads are not pinned to specific nodes.
7    Disabled,
8    /// Distributes threads across NUMA nodes in a round-robin fashion.
9    Distribute,
10    /// Pins all threads to the current NUMA node to avoid cross-node memory access.
11    Isolate,
12    /// Respects the CPU affinity mask set externally by the `numactl` command.
13    Numactl,
14    /// Mirrors memory across NUMA nodes. Currently a no-op in llama.cpp.
15    Mirror,
16}
17
18impl TryFrom<llama_cpp_bindings_sys::ggml_numa_strategy> for NumaStrategy {
19    type Error = InvalidNumaStrategy;
20
21    fn try_from(value: llama_cpp_bindings_sys::ggml_numa_strategy) -> Result<Self, Self::Error> {
22        match value {
23            llama_cpp_bindings_sys::GGML_NUMA_STRATEGY_DISABLED => Ok(Self::Disabled),
24            llama_cpp_bindings_sys::GGML_NUMA_STRATEGY_DISTRIBUTE => Ok(Self::Distribute),
25            llama_cpp_bindings_sys::GGML_NUMA_STRATEGY_ISOLATE => Ok(Self::Isolate),
26            llama_cpp_bindings_sys::GGML_NUMA_STRATEGY_NUMACTL => Ok(Self::Numactl),
27            llama_cpp_bindings_sys::GGML_NUMA_STRATEGY_MIRROR => Ok(Self::Mirror),
28            value => Err(InvalidNumaStrategy(value)),
29        }
30    }
31}
32
33impl From<NumaStrategy> for llama_cpp_bindings_sys::ggml_numa_strategy {
34    fn from(value: NumaStrategy) -> Self {
35        match value {
36            NumaStrategy::Disabled => llama_cpp_bindings_sys::GGML_NUMA_STRATEGY_DISABLED,
37            NumaStrategy::Distribute => llama_cpp_bindings_sys::GGML_NUMA_STRATEGY_DISTRIBUTE,
38            NumaStrategy::Isolate => llama_cpp_bindings_sys::GGML_NUMA_STRATEGY_ISOLATE,
39            NumaStrategy::Numactl => llama_cpp_bindings_sys::GGML_NUMA_STRATEGY_NUMACTL,
40            NumaStrategy::Mirror => llama_cpp_bindings_sys::GGML_NUMA_STRATEGY_MIRROR,
41        }
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::NumaStrategy;
48    use crate::invalid_numa_strategy::InvalidNumaStrategy;
49
50    #[test]
51    fn numa_from_and_to() {
52        let numas = [
53            NumaStrategy::Disabled,
54            NumaStrategy::Distribute,
55            NumaStrategy::Isolate,
56            NumaStrategy::Numactl,
57            NumaStrategy::Mirror,
58        ];
59
60        for numa in &numas {
61            let raw = llama_cpp_bindings_sys::ggml_numa_strategy::from(*numa);
62            let roundtripped =
63                NumaStrategy::try_from(raw).expect("Failed to roundtrip NumaStrategy");
64
65            assert_eq!(*numa, roundtripped);
66        }
67    }
68
69    #[test]
70    fn invalid_numa_strategy_returns_error() {
71        let invalid_value = 800;
72        let result = NumaStrategy::try_from(invalid_value);
73
74        assert_eq!(result, Err(InvalidNumaStrategy(invalid_value)));
75    }
76}