fib_quant/
ffi.rs

1//! FFI bindings to the C kernels in `c-kernels/`.
2//!
3//! These wrappers provide safe Rust interfaces to the compiled C
4//! implementations of the hot numerical loops. The Rust side handles
5//! all validation; the C code only performs the tight inner loops.
6
7// This module requires unsafe FFI calls to the C kernels. The crate-level
8// `unsafe_code = "deny"` lint is overridden here because FFI is the
9// intended and reviewed use of unsafe in this module.
10#![allow(unsafe_code)]
11
12use std::os::raw::c_int;
13
14extern "C" {
15    fn fq_encode_vector_block(
16        vec: *const f32,
17        dim: usize,
18        codebook: *const f32,
19        codebook_size: usize,
20        block_dim: usize,
21        out_indices: *mut u16,
22    ) -> usize;
23
24    fn fq_decode_vector_block(
25        indices: *const u16,
26        block_count: usize,
27        codebook: *const f32,
28        codebook_size: usize,
29        block_dim: usize,
30        out_vec: *mut f32,
31    );
32
33    fn fq_compressed_attention_logits(
34        key_indices: *const u16,
35        n_keys: usize,
36        key_norms: *const f32,
37        gram_table: *const f32,
38        gram_size: usize,
39        query_indices: *const u16,
40        query_count: usize,
41        query_norm: f32,
42        scale: f32,
43        out_logits: *mut f32,
44    );
45
46    fn fq_softmax(logits: *mut f32, n: usize) -> c_int;
47}
48
49/// Encode a vector block: find nearest codeword for each block of
50/// `block_dim` floats. Returns the codeword indices.
51///
52/// # Safety
53/// The caller must ensure:
54/// - `vec.len() == dim` and `dim % block_dim == 0`
55/// - `codebook.len() == codebook_size * block_dim`
56/// - `out_indices` has capacity `dim / block_dim`
57pub fn c_encode_vector_block(
58    vec: &[f32],
59    codebook: &[f32],
60    codebook_size: usize,
61    block_dim: usize,
62) -> Vec<u16> {
63    let dim = vec.len();
64    let block_count = dim / block_dim;
65    let mut indices = vec![0u16; block_count];
66    // SAFETY: All pointers are valid slices with correct lengths.
67    // The C function reads `dim` floats from vec, `codebook_size*block_dim`
68    // floats from codebook, and writes `block_count` u16 to out_indices.
69    unsafe {
70        fq_encode_vector_block(
71            vec.as_ptr(),
72            dim,
73            codebook.as_ptr(),
74            codebook_size,
75            block_dim,
76            indices.as_mut_ptr(),
77        );
78    }
79    indices
80}
81
82/// Decode a vector block: gather codewords by index.
83///
84/// # Panics
85/// Panics if any index >= codebook_size (caller should validate).
86pub fn c_decode_vector_block(
87    indices: &[u16],
88    codebook: &[f32],
89    codebook_size: usize,
90    block_dim: usize,
91) -> Vec<f32> {
92    let block_count = indices.len();
93    let mut out = vec![0.0f32; block_count * block_dim];
94    // SAFETY: All pointers are valid slices with correct lengths.
95    unsafe {
96        fq_decode_vector_block(
97            indices.as_ptr(),
98            block_count,
99            codebook.as_ptr(),
100            codebook_size,
101            block_dim,
102            out.as_mut_ptr(),
103        );
104    }
105    out
106}
107
108/// Compute compressed attention logits via C kernel.
109///
110/// For each key, sums Gram table lookups across all blocks, then scales
111/// by query_norm * stored_norm / scale.
112#[allow(clippy::too_many_arguments)]
113pub fn c_compressed_attention_logits(
114    key_indices: &[u16],
115    n_keys: usize,
116    key_norms: &[f32],
117    gram_table: &[f32],
118    gram_size: usize,
119    query_indices: &[u16],
120    query_count: usize,
121    query_norm: f32,
122    scale: f32,
123) -> Vec<f32> {
124    let mut logits = vec![0.0f32; n_keys];
125    // SAFETY: All pointers are valid slices with correct lengths.
126    unsafe {
127        fq_compressed_attention_logits(
128            key_indices.as_ptr(),
129            n_keys,
130            key_norms.as_ptr(),
131            gram_table.as_ptr(),
132            gram_size,
133            query_indices.as_ptr(),
134            query_count,
135            query_norm,
136            scale,
137            logits.as_mut_ptr(),
138        );
139    }
140    logits
141}
142
143/// Numerically stable softmax via C kernel.
144///
145/// Returns `Ok(())` on success, `Err` on underflow or empty input.
146#[allow(clippy::result_unit_err)]
147pub fn c_softmax(logits: &mut [f32]) -> Result<(), ()> {
148    if logits.is_empty() {
149        return Err(());
150    }
151    // SAFETY: logits is a valid mutable slice.
152    let ret = unsafe { fq_softmax(logits.as_mut_ptr(), logits.len()) };
153    if ret == 0 {
154        Ok(())
155    } else {
156        Err(())
157    }
158}
fib_quant/ffi.rs

fib_quant/
ffi.rs
