diff --git a/crates/goose-local-inference/src/llamacpp/mod.rs b/crates/goose-local-inference/src/llamacpp/mod.rs index bdfd8cf41e10..9b7baba22549 100644 --- a/crates/goose-local-inference/src/llamacpp/mod.rs +++ b/crates/goose-local-inference/src/llamacpp/mod.rs @@ -303,12 +303,51 @@ fn select_generation_template<'a>( }) } +#[cfg(any(target_arch = "x86_64", test))] +fn unsupported_cpu_features_error_message(missing_features: &[&str]) -> String { + format!( + "Local inference with the bundled llama.cpp backend requires CPU support for {}. \ + This CPU is missing {}. Use a CPU with those instruction sets or switch to a non-local provider.", + missing_features.join(" and "), + missing_features.join(", ") + ) +} + +#[cfg(target_arch = "x86_64")] +fn check_cpu_supports_local_inference() -> Result<()> { + let missing_features = [ + (!std::arch::is_x86_feature_detected!("fma")).then_some("FMA"), + (!std::arch::is_x86_feature_detected!("avx2")).then_some("AVX2"), + (!std::arch::is_x86_feature_detected!("f16c")).then_some("F16C"), + (!std::arch::is_x86_feature_detected!("bmi2")).then_some("BMI2"), + (!std::arch::is_x86_feature_detected!("sse4.2")).then_some("SSE4.2"), + ] + .into_iter() + .flatten() + .collect::>(); + + if missing_features.is_empty() { + Ok(()) + } else { + Err(anyhow::anyhow!(unsupported_cpu_features_error_message( + &missing_features + ))) + } +} + +#[cfg(not(target_arch = "x86_64"))] +fn check_cpu_supports_local_inference() -> Result<()> { + Ok(()) +} + pub(super) struct LlamaCppBackend { backend: LlamaBackend, } impl LlamaCppBackend { pub(super) fn new() -> Result { + check_cpu_supports_local_inference()?; + let backend = match LlamaBackend::init() { Ok(backend) => backend, Err(llama_cpp_2::LlamaCppError::BackendAlreadyInitialized) => { @@ -703,4 +742,38 @@ mod tests { "{% for message in messages %}{{ message.content }}{% endfor %}" )); } + + #[test] + fn local_inference_cpu_support_check_accepts_current_host() { + let result = check_cpu_supports_local_inference(); + + #[cfg(target_arch = "x86_64")] + { + let supports_required_features = std::arch::is_x86_feature_detected!("fma") + && std::arch::is_x86_feature_detected!("avx2") + && std::arch::is_x86_feature_detected!("f16c") + && std::arch::is_x86_feature_detected!("bmi2") + && std::arch::is_x86_feature_detected!("sse4.2"); + + if supports_required_features { + assert!( + result.is_ok(), + "expected supported x86_64 host to pass: {result:?}" + ); + } else { + assert!(result.is_err(), "expected unsupported x86_64 host to fail"); + } + } + + #[cfg(not(target_arch = "x86_64"))] + assert!(result.is_ok()); + } + + #[test] + fn local_inference_cpu_support_error_names_missing_instruction_sets() { + let message = unsupported_cpu_features_error_message(&["FMA", "AVX2"]); + + assert!(message.contains("FMA")); + assert!(message.contains("AVX2")); + } }