ailake_index/
lib.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! ailake-index — HNSW and IVF-PQ index lifecycle
3//!
4//! Search backend priority (all detected at runtime, no build-time GPU SDK required):
5//!   1. NVIDIA CUDA — cuBLAS SGEMM via libloading dlopen of libcudart + libcublas
6//!   2. AMD ROCm    — hipBLAS SGEMM via libloading dlopen of libamdhip64 + libhipblas
7//!   3. CPU rayon   — parallel brute-force, always available
8
9pub mod gpu;
10pub mod hardware;
11pub mod hnsw;
12pub mod ivf_pq;
13pub mod mmap_loader;
14pub mod serialize;
15
16pub use hardware::{detect_backend, detect_cuda, detect_rocm, HardwareBackend, HardwareProfile};
17pub use hnsw::{HnswBuilder, HnswConfig, HnswIndex};
18pub use ivf_pq::{find_valid_pq_m, IvfPqCodebook, IvfPqConfig, IvfPqIndex, IvfPqSerializer};
19pub use mmap_loader::MmapLoader;
20pub use serialize::HnswSerializer;
21
22use ailake_core::RowId;
23
24/// Unified index type: dispatches search to HNSW or IVF-PQ.
25pub enum AnyIndex {
26    Hnsw(HnswIndex),
27    IvfPq(IvfPqIndex),
28}
29
30impl AnyIndex {
31    /// Search. `ef` is used for HNSW; ignored for IVF-PQ.
32    pub fn search(&self, query: &[f32], top_k: usize, ef: usize) -> Vec<(RowId, f32)> {
33        match self {
34            AnyIndex::Hnsw(idx) => idx.search(query, top_k, ef),
35            AnyIndex::IvfPq(idx) => idx.search(query, top_k, None),
36        }
37    }
38
39    pub fn node_count(&self) -> u64 {
40        match self {
41            AnyIndex::Hnsw(idx) => idx.node_count(),
42            AnyIndex::IvfPq(idx) => idx.node_count(),
43        }
44    }
45
46    /// Quantize stored vectors to F16 for HNSW search (no-op for IVF-PQ).
47    pub fn quantize_to_f16(&mut self) {
48        if let AnyIndex::Hnsw(idx) = self {
49            idx.quantize_to_f16();
50        }
51    }
52
53    pub fn is_hnsw(&self) -> bool {
54        matches!(self, AnyIndex::Hnsw(_))
55    }
56
57    pub fn is_ivf_pq(&self) -> bool {
58        matches!(self, AnyIndex::IvfPq(_))
59    }
60}
ailake_index/lib.rs

ailake_index/
lib.rs
