Spaces:
Running
Running
File size: 1,832 Bytes
ba86059 893e0bb ba86059 893e0bb ba86059 f267654 ba86059 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | use pyo3::prelude::*;
#[cfg(target_os = "windows")]
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
mod turboquant;
mod sq8;
mod pq;
mod baselines;
#[pymodule]
fn tq_native_lib(m: &Bound<'_, PyModule>) -> PyResult<()> {
// Threading:
// - Default: use OS-reported available parallelism
// - Override: set env var TQ_RAYON_THREADS (e.g. "8")
// NOTE: build_global() can only be called once; ignore error if already initialized.
let threads = 2* std::env::var("TQ_RAYON_THREADS")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.or_else(|| std::thread::available_parallelism().ok().map(|n| n.get()))
.unwrap_or(0);
let mut b = rayon::ThreadPoolBuilder::new();
if threads > 0 {
b = b.num_threads(threads);
}
let _ = b.build_global();
// NEW CORE: Scalar Quantization (b-1) + QJL (1)
m.add_function(wrap_pyfunction!(turboquant::tq_scan, m)?)?;
m.add_function(wrap_pyfunction!(turboquant::tq_batch_scan, m)?)?;
m.add_function(wrap_pyfunction!(turboquant::tq_ivf_online_scan, m)?)?;
m.add_function(wrap_pyfunction!(turboquant::tq_ivf_scan_with_clusters, m)?)?;
m.add_function(wrap_pyfunction!(turboquant::tq_quantize_rotated, m)?)?;
m.add_function(wrap_pyfunction!(turboquant::tq_kmeans_train, m)?)?;
m.add_function(wrap_pyfunction!(turboquant::tq_assign_clusters, m)?)?;
m.add_function(wrap_pyfunction!(turboquant::tq_unified_search, m)?)?;
// Baseline comparisons (New names for stress_5m.py)
m.add_function(wrap_pyfunction!(baselines::sq_scan, m)?)?;
m.add_function(wrap_pyfunction!(baselines::pq_scan, m)?)?;
// Original names
m.add_function(wrap_pyfunction!(sq8::sq8_score_simd, m)?)?;
m.add_function(wrap_pyfunction!(pq::pq_score_simd, m)?)?;
Ok(())
}
|