Initial ForgeLM import

crates/llama-sys/src/lib.rs

@@ -0,0 +1,571 @@
+//! llama-sys: Dynamic FFI bindings to llama.cpp via libloading.
+//!
+//! Loads `llama.dll` (Windows) / `libllama.so` (Linux) / `libllama.dylib` (macOS)
+//! at runtime from the path specified by `LLAMA_DLL_PATH` env var, or from a
+//! set of well-known locations.
+//!
+//! This approach requires no `.lib` import library and no source compilation.
+
+#![allow(non_upper_case_globals)]
+#![allow(non_camel_case_types)]
+#![allow(non_snake_case)]
+#![allow(dead_code)]
+#![allow(clippy::all)]
+
+use std::ffi::{c_char, c_void};
+use std::path::{Path, PathBuf};
+
+pub use libloading::Library;
+
+// ─── Primitive type aliases ────────────────────────────────────────────────
+
+pub type llama_pos = i32;
+pub type llama_token = i32;
+pub type llama_seq_id = i32;
+
+// ─── Opaque struct handles ─────────────────────────────────────────────────
+
+#[repr(C)]
+pub struct llama_model {
+    _private: [u8; 0],
+}
+
+#[repr(C)]
+pub struct llama_context {
+    _private: [u8; 0],
+}
+
+#[repr(C)]
+pub struct llama_vocab {
+    _private: [u8; 0],
+}
+
+#[repr(C)]
+pub struct llama_sampler {
+    _private: [u8; 0],
+}
+
+#[repr(C)]
+pub struct llama_memory_i {
+    _private: [u8; 0],
+}
+
+pub type llama_memory_t = *mut llama_memory_i;
+
+// ─── Model params ──────────────────────────────────────────────────────────
+
+#[repr(C)]
+pub struct llama_model_params {
+    pub devices: *mut c_void,
+    pub tensor_buft_overrides: *const c_void,
+    pub n_gpu_layers: i32,
+    pub split_mode: i32,
+    pub main_gpu: i32,
+    pub tensor_split: *const f32,
+    pub progress_callback: Option<extern "C" fn(f32, *mut c_void) -> bool>,
+    pub progress_callback_user_data: *mut c_void,
+    pub kv_overrides: *const c_void,
+    pub vocab_only: bool,
+    pub use_mmap: bool,
+    pub use_direct_io: bool,
+    pub use_mlock: bool,
+    pub check_tensors: bool,
+    pub use_extra_bufts: bool,
+    pub no_host: bool,
+    pub no_alloc: bool,
+}
+
+// ─── Context params ────────────────────────────────────────────────────────
+
+#[repr(C)]
+pub struct llama_context_params {
+    pub n_ctx: u32,
+    pub n_batch: u32,
+    pub n_ubatch: u32,
+    pub n_seq_max: u32,
+    pub n_threads: i32,
+    pub n_threads_batch: i32,
+    pub rope_scaling_type: i32,
+    pub pooling_type: i32,
+    pub attention_type: i32,
+    pub flash_attn_type: i32,
+    pub rope_freq_base: f32,
+    pub rope_freq_scale: f32,
+    pub yarn_ext_factor: f32,
+    pub yarn_attn_factor: f32,
+    pub yarn_beta_fast: f32,
+    pub yarn_beta_slow: f32,
+    pub yarn_orig_ctx: u32,
+    pub defrag_thold: f32,
+    pub cb_eval: Option<extern "C" fn(*mut c_void, bool) -> bool>,
+    pub cb_eval_user_data: *mut c_void,
+    pub type_k: i32,
+    pub type_v: i32,
+    pub abort_callback: Option<extern "C" fn(*mut c_void) -> bool>,
+    pub abort_callback_data: *mut c_void,
+    pub embeddings: bool,
+    pub offload_kqv: bool,
+    pub no_perf: bool,
+    pub op_offload: bool,
+    pub swa_full: bool,
+    pub kv_unified: bool,
+    pub samplers: *mut c_void,
+    pub n_samplers: usize,
+}
+
+// ─── Batch ─────────────────────────────────────────────────────────────────
+
+#[repr(C)]
+pub struct llama_batch {
+    pub n_tokens: i32,
+    pub token: *mut llama_token,
+    pub embd: *mut f32,
+    pub pos: *mut llama_pos,
+    pub n_seq_id: *mut i32,
+    pub seq_id: *mut *mut llama_seq_id,
+    pub logits: *mut i8,
+}
+
+// ─── Sampler chain params ──────────────────────────────────────────────────
+
+#[repr(C)]
+pub struct llama_sampler_chain_params {
+    pub no_perf: bool,
+}
+
+// ─── Flash attn type ───────────────────────────────────────────────────────
+
+pub const LLAMA_FLASH_ATTN_TYPE_AUTO: i32 = -1;
+pub const LLAMA_FLASH_ATTN_TYPE_DISABLED: i32 = 0;
+pub const LLAMA_FLASH_ATTN_TYPE_ENABLED: i32 = 1;
+
+// ─── Dynamic loader ────────────────────────────────────────────────────────
+
+/// Container for all resolved function pointers from llama.dll.
+pub struct LlamaLib {
+    // Keep the library alive as long as this struct is alive.
+    _lib: Library,
+
+    // ── Lifecycle ──────────────────────────────────────────────────────────
+    pub backend_init: unsafe extern "C" fn(),
+    pub backend_free: unsafe extern "C" fn(),
+    pub backend_load_all: unsafe extern "C" fn(),
+
+    // ── Defaults ──────────────────────────────────────────────────────────
+    pub model_default_params: unsafe extern "C" fn() -> llama_model_params,
+    pub context_default_params: unsafe extern "C" fn() -> llama_context_params,
+    pub sampler_chain_default_params: unsafe extern "C" fn() -> llama_sampler_chain_params,
+
+    // ── Model ─────────────────────────────────────────────────────────────
+    pub model_load_from_file:
+        unsafe extern "C" fn(*const c_char, llama_model_params) -> *mut llama_model,
+    pub model_free: unsafe extern "C" fn(*mut llama_model),
+    pub model_get_vocab: unsafe extern "C" fn(*const llama_model) -> *const llama_vocab,
+    pub model_n_ctx_train: unsafe extern "C" fn(*const llama_model) -> i32,
+    pub model_is_recurrent: unsafe extern "C" fn(*const llama_model) -> bool,
+    pub model_is_hybrid: unsafe extern "C" fn(*const llama_model) -> bool,
+    pub model_desc: unsafe extern "C" fn(*const llama_model, *mut c_char, usize) -> i32,
+    pub model_n_params: unsafe extern "C" fn(*const llama_model) -> u64,
+    pub model_size: unsafe extern "C" fn(*const llama_model) -> u64,
+
+    // ── Context ───────────────────────────────────────────────────────────
+    pub init_from_model:
+        unsafe extern "C" fn(*mut llama_model, llama_context_params) -> *mut llama_context,
+    pub free: unsafe extern "C" fn(*mut llama_context),
+    pub n_ctx: unsafe extern "C" fn(*const llama_context) -> u32,
+    pub get_memory: unsafe extern "C" fn(*const llama_context) -> llama_memory_t,
+
+    // ── Batch ─────────────────────────────────────────────────────────────
+    pub batch_get_one: unsafe extern "C" fn(*mut llama_token, i32) -> llama_batch,
+    pub batch_init: unsafe extern "C" fn(i32, i32, i32) -> llama_batch,
+    pub batch_free: unsafe extern "C" fn(llama_batch),
+
+    // ── Decode ────────────────────────────────────────────────────────────
+    pub decode: unsafe extern "C" fn(*mut llama_context, llama_batch) -> i32,
+    pub set_n_threads: unsafe extern "C" fn(*mut llama_context, i32, i32),
+    pub set_warmup: unsafe extern "C" fn(*mut llama_context, bool),
+    pub synchronize: unsafe extern "C" fn(*mut llama_context),
+    pub memory_clear: unsafe extern "C" fn(llama_memory_t, bool),
+
+    // ── Logits ────────────────────────────────────────────────────────────
+    pub get_logits_ith: unsafe extern "C" fn(*mut llama_context, i32) -> *mut f32,
+
+    // ── Vocab / Tokenizer ─────────────────────────────────────────────────
+    pub vocab_n_tokens: unsafe extern "C" fn(*const llama_vocab) -> i32,
+    pub vocab_is_eog: unsafe extern "C" fn(*const llama_vocab, llama_token) -> bool,
+    pub vocab_get_text: unsafe extern "C" fn(*const llama_vocab, llama_token) -> *const c_char,
+    pub token_to_piece: Option<
+        unsafe extern "C" fn(*const llama_vocab, llama_token, *mut c_char, i32, i32, bool) -> i32,
+    >,
+    pub tokenize: unsafe extern "C" fn(
+        *const llama_vocab,
+        *const c_char,
+        i32,
+        *mut llama_token,
+        i32,
+        bool,
+        bool,
+    ) -> i32,
+    pub token_bos: unsafe extern "C" fn(*const llama_vocab) -> llama_token,
+    pub token_eos: unsafe extern "C" fn(*const llama_vocab) -> llama_token,
+
+    // ── Sampler ───────────────────────────────────────────────────────────
+    pub sampler_chain_init: unsafe extern "C" fn(llama_sampler_chain_params) -> *mut llama_sampler,
+    pub sampler_chain_add: unsafe extern "C" fn(*mut llama_sampler, *mut llama_sampler),
+    pub sampler_free: unsafe extern "C" fn(*mut llama_sampler),
+    pub sampler_sample:
+        unsafe extern "C" fn(*mut llama_sampler, *mut llama_context, i32) -> llama_token,
+    pub sampler_init_greedy: unsafe extern "C" fn() -> *mut llama_sampler,
+    pub sampler_init_temp: unsafe extern "C" fn(f32) -> *mut llama_sampler,
+    pub sampler_init_top_p: unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
+    pub sampler_init_top_k: unsafe extern "C" fn(i32) -> *mut llama_sampler,
+    pub sampler_init_min_p: unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
+    pub sampler_init_dist: unsafe extern "C" fn(u32) -> *mut llama_sampler,
+
+    // ── Misc ──────────────────────────────────────────────────────────────
+    pub supports_gpu_offload: unsafe extern "C" fn() -> bool,
+    pub time_us: unsafe extern "C" fn() -> i64,
+}
+
+impl LlamaLib {
+    /// Load the llama shared library from the given path.
+    pub fn load(path: &Path) -> anyhow::Result<Self> {
+        use anyhow::Context;
+
+        let lib = unsafe { Library::new(path) }
+            .with_context(|| format!("Failed to load llama library from {:?}", path))?;
+
+        macro_rules! sym {
+            ($name:literal) => {{
+                *unsafe { lib.get::<unsafe extern "C" fn()>(concat!($name, "\0").as_bytes()) }
+                    .with_context(|| format!("Missing symbol: {}", $name))?
+            }};
+        }
+
+        // We need transmutes to specific fn types — use a helper macro:
+        macro_rules! load {
+            ($field:ident, $ty:ty, $name:literal) => {{
+                let raw: libloading::Symbol<unsafe extern "C" fn()> = unsafe {
+                    lib.get(concat!($name, "\0").as_bytes())
+                        .with_context(|| format!("Missing symbol: {}", $name))?
+                };
+                unsafe { std::mem::transmute::<_, $ty>(*raw) }
+            }};
+        }
+
+        Ok(Self {
+            backend_init: load!(backend_init, unsafe extern "C" fn(), "llama_backend_init"),
+            backend_free: load!(backend_free, unsafe extern "C" fn(), "llama_backend_free"),
+            backend_load_all: {
+                let sym_name = b"ggml_backend_load_all\0";
+                if let Ok(raw) = unsafe { lib.get::<unsafe extern "C" fn()>(sym_name) } {
+                    unsafe { std::mem::transmute(*raw) }
+                } else {
+                    let ggml_path = path.with_file_name("ggml.dll");
+                    if let Ok(ggml_lib) = unsafe { libloading::Library::new(&ggml_path) } {
+                        // Faz a extração do raw *antes* de dar leak no Box, desfazendo ownership borrow issue do Rust
+                        let raw_ptr = if let Ok(raw) =
+                            unsafe { ggml_lib.get::<unsafe extern "C" fn()>(sym_name) }
+                        {
+                            Some(unsafe { std::mem::transmute::<_, unsafe extern "C" fn()>(*raw) })
+                        } else {
+                            None
+                        };
+
+                        // Faz leak no handle do módulo pra manter vivo
+                        std::mem::forget(ggml_lib);
+
+                        raw_ptr.unwrap_or_else(|| {
+                            panic!("Missing symbol ggml_backend_load_all in both llama.dll and ggml.dll")
+                        })
+                    } else {
+                        anyhow::bail!("Missing symbol ggml_backend_load_all in llama.dll and could not load ggml.dll");
+                    }
+                }
+            },
+
+            model_default_params: load!(
+                model_default_params,
+                unsafe extern "C" fn() -> llama_model_params,
+                "llama_model_default_params"
+            ),
+            context_default_params: load!(
+                context_default_params,
+                unsafe extern "C" fn() -> llama_context_params,
+                "llama_context_default_params"
+            ),
+            sampler_chain_default_params: load!(
+                sampler_chain_default_params,
+                unsafe extern "C" fn() -> llama_sampler_chain_params,
+                "llama_sampler_chain_default_params"
+            ),
+
+            model_load_from_file: load!(
+                model_load_from_file,
+                unsafe extern "C" fn(*const c_char, llama_model_params) -> *mut llama_model,
+                "llama_model_load_from_file"
+            ),
+            model_free: load!(
+                model_free,
+                unsafe extern "C" fn(*mut llama_model),
+                "llama_model_free"
+            ),
+            model_get_vocab: load!(
+                model_get_vocab,
+                unsafe extern "C" fn(*const llama_model) -> *const llama_vocab,
+                "llama_model_get_vocab"
+            ),
+            model_n_ctx_train: load!(
+                model_n_ctx_train,
+                unsafe extern "C" fn(*const llama_model) -> i32,
+                "llama_model_n_ctx_train"
+            ),
+            model_is_recurrent: load!(
+                model_is_recurrent,
+                unsafe extern "C" fn(*const llama_model) -> bool,
+                "llama_model_is_recurrent"
+            ),
+            model_is_hybrid: load!(
+                model_is_hybrid,
+                unsafe extern "C" fn(*const llama_model) -> bool,
+                "llama_model_is_hybrid"
+            ),
+            model_desc: load!(
+                model_desc,
+                unsafe extern "C" fn(*const llama_model, *mut c_char, usize) -> i32,
+                "llama_model_desc"
+            ),
+            model_n_params: load!(
+                model_n_params,
+                unsafe extern "C" fn(*const llama_model) -> u64,
+                "llama_model_n_params"
+            ),
+            model_size: load!(
+                model_size,
+                unsafe extern "C" fn(*const llama_model) -> u64,
+                "llama_model_size"
+            ),
+
+            init_from_model: load!(
+                init_from_model,
+                unsafe extern "C" fn(*mut llama_model, llama_context_params) -> *mut llama_context,
+                "llama_init_from_model"
+            ),
+            free: load!(free, unsafe extern "C" fn(*mut llama_context), "llama_free"),
+            n_ctx: load!(
+                n_ctx,
+                unsafe extern "C" fn(*const llama_context) -> u32,
+                "llama_n_ctx"
+            ),
+            get_memory: load!(
+                get_memory,
+                unsafe extern "C" fn(*const llama_context) -> llama_memory_t,
+                "llama_get_memory"
+            ),
+
+            batch_get_one: load!(
+                batch_get_one,
+                unsafe extern "C" fn(*mut llama_token, i32) -> llama_batch,
+                "llama_batch_get_one"
+            ),
+            batch_init: load!(
+                batch_init,
+                unsafe extern "C" fn(i32, i32, i32) -> llama_batch,
+                "llama_batch_init"
+            ),
+            batch_free: load!(
+                batch_free,
+                unsafe extern "C" fn(llama_batch),
+                "llama_batch_free"
+            ),
+
+            decode: load!(
+                decode,
+                unsafe extern "C" fn(*mut llama_context, llama_batch) -> i32,
+                "llama_decode"
+            ),
+            set_n_threads: load!(
+                set_n_threads,
+                unsafe extern "C" fn(*mut llama_context, i32, i32),
+                "llama_set_n_threads"
+            ),
+            set_warmup: load!(
+                set_warmup,
+                unsafe extern "C" fn(*mut llama_context, bool),
+                "llama_set_warmup"
+            ),
+            synchronize: load!(
+                synchronize,
+                unsafe extern "C" fn(*mut llama_context),
+                "llama_synchronize"
+            ),
+            memory_clear: load!(
+                memory_clear,
+                unsafe extern "C" fn(llama_memory_t, bool),
+                "llama_memory_clear"
+            ),
+
+            get_logits_ith: load!(
+                get_logits_ith,
+                unsafe extern "C" fn(*mut llama_context, i32) -> *mut f32,
+                "llama_get_logits_ith"
+            ),
+
+            vocab_n_tokens: load!(
+                vocab_n_tokens,
+                unsafe extern "C" fn(*const llama_vocab) -> i32,
+                "llama_vocab_n_tokens"
+            ),
+            vocab_is_eog: load!(
+                vocab_is_eog,
+                unsafe extern "C" fn(*const llama_vocab, llama_token) -> bool,
+                "llama_vocab_is_eog"
+            ),
+            vocab_get_text: load!(
+                vocab_get_text,
+                unsafe extern "C" fn(*const llama_vocab, llama_token) -> *const c_char,
+                "llama_vocab_get_text"
+            ),
+            token_to_piece: unsafe {
+                lib.get::<unsafe extern "C" fn(
+                    *const llama_vocab,
+                    llama_token,
+                    *mut c_char,
+                    i32,
+                    i32,
+                    bool,
+                ) -> i32>(b"llama_token_to_piece\0")
+                    .ok()
+                    .map(|raw| *raw)
+            },
+            tokenize: load!(
+                tokenize,
+                unsafe extern "C" fn(
+                    *const llama_vocab,
+                    *const c_char,
+                    i32,
+                    *mut llama_token,
+                    i32,
+                    bool,
+                    bool,
+                ) -> i32,
+                "llama_tokenize"
+            ),
+            token_bos: load!(
+                token_bos,
+                unsafe extern "C" fn(*const llama_vocab) -> llama_token,
+                "llama_vocab_bos"
+            ),
+            token_eos: load!(
+                token_eos,
+                unsafe extern "C" fn(*const llama_vocab) -> llama_token,
+                "llama_vocab_eos"
+            ),
+
+            sampler_chain_init: load!(
+                sampler_chain_init,
+                unsafe extern "C" fn(llama_sampler_chain_params) -> *mut llama_sampler,
+                "llama_sampler_chain_init"
+            ),
+            sampler_chain_add: load!(
+                sampler_chain_add,
+                unsafe extern "C" fn(*mut llama_sampler, *mut llama_sampler),
+                "llama_sampler_chain_add"
+            ),
+            sampler_free: load!(
+                sampler_free,
+                unsafe extern "C" fn(*mut llama_sampler),
+                "llama_sampler_free"
+            ),
+            sampler_sample: load!(
+                sampler_sample,
+                unsafe extern "C" fn(*mut llama_sampler, *mut llama_context, i32) -> llama_token,
+                "llama_sampler_sample"
+            ),
+            sampler_init_greedy: load!(
+                sampler_init_greedy,
+                unsafe extern "C" fn() -> *mut llama_sampler,
+                "llama_sampler_init_greedy"
+            ),
+            sampler_init_temp: load!(
+                sampler_init_temp,
+                unsafe extern "C" fn(f32) -> *mut llama_sampler,
+                "llama_sampler_init_temp"
+            ),
+            sampler_init_top_p: load!(
+                sampler_init_top_p,
+                unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
+                "llama_sampler_init_top_p"
+            ),
+            sampler_init_top_k: load!(
+                sampler_init_top_k,
+                unsafe extern "C" fn(i32) -> *mut llama_sampler,
+                "llama_sampler_init_top_k"
+            ),
+            sampler_init_min_p: load!(
+                sampler_init_min_p,
+                unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
+                "llama_sampler_init_min_p"
+            ),
+            sampler_init_dist: load!(
+                sampler_init_dist,
+                unsafe extern "C" fn(u32) -> *mut llama_sampler,
+                "llama_sampler_init_dist"
+            ),
+
+            supports_gpu_offload: load!(
+                supports_gpu_offload,
+                unsafe extern "C" fn() -> bool,
+                "llama_supports_gpu_offload"
+            ),
+            time_us: load!(time_us, unsafe extern "C" fn() -> i64, "llama_time_us"),
+
+            _lib: lib,
+        })
+    }
+
+    /// Resolve the DLL path from env var or known locations.
+    pub fn find_dll() -> anyhow::Result<PathBuf> {
+        // 1. Explicit env override
+        if let Ok(p) = std::env::var("LLAMA_DLL_PATH") {
+            let path = PathBuf::from(p);
+            if path.exists() {
+                return Ok(path);
+            }
+        }
+
+        // 2. Load directly relying on the Windows Library loader searching the .exe dir
+        // since our build.rs now copies all DLLs to target/{debug,release}/
+        #[cfg(target_os = "windows")]
+        {
+            return Ok(PathBuf::from("llama.dll"));
+        }
+
+        // 3. Unix-like: rely on LD_LIBRARY_PATH / rpath
+        #[cfg(not(target_os = "windows"))]
+        {
+            for name in &["libllama.so", "libllama.dylib"] {
+                if let Ok(lib) = Library::new(name) {
+                    drop(lib);
+                    return Ok(PathBuf::from(name));
+                }
+            }
+        }
+
+        anyhow::bail!(
+            "Cannot find llama shared library. Set LLAMA_DLL_PATH to the full path of llama.dll"
+        )
+    }
+
+    /// Convenience: find + load in one call.
+    pub fn load_auto() -> anyhow::Result<Self> {
+        let path = Self::find_dll()?;
+        tracing::info!(path = %path.display(), "Loading llama library");
+        Self::load(&path)
+    }
+}
+
+// SAFETY: LlamaLib only holds a Library and function pointers.
+// All calls must be externally synchronized (enforced by InferenceEngine).
+unsafe impl Send for LlamaLib {}
+unsafe impl Sync for LlamaLib {}