hipvs/ivf_flat/
mod.rs

1/*
2 * Copyright (c) 2024, NVIDIA CORPORATION.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! The IVF-Flat method is an ANN algorithm. It uses an inverted file index (IVF) with
18//! unmodified (that is, flat) vectors. This algorithm provides simple knobs to reduce
19//! the overall search space and to trade-off accuracy for speed.
20//!
21//! Example:
22//! ```
23//!
24//! use hipvs::ivf_flat::{Index, IndexParams, SearchParams};
25//! use hipvs::{ManagedTensor, Resources, Result};
26//!
27//! use ndarray::s;
28//! use ndarray_rand::rand_distr::Uniform;
29//! use ndarray_rand::RandomExt;
30//!
31//! fn ivf_flat_example() -> Result<()> {
32//!     let res = Resources::new()?;
33//!
34//!     // Create a new random dataset to index
35//!     let n_datapoints = 65536;
36//!     let n_features = 512;
37//!     let dataset =
38//!         ndarray::Array::<f32, _>::random((n_datapoints, n_features), Uniform::new(0., 1.0));
39//!
40//!     // build the ivf-flat index
41//!     let build_params = IndexParams::new()?;
42//!     let index = Index::build(&res, &build_params, &dataset)?;
43//!     println!(
44//!         "Indexed {}x{} datapoints into ivf-flat index",
45//!         n_datapoints, n_features
46//!     );
47//!
48//!     // use the first 4 points from the dataset as queries : will test that we get them back
49//!     // as their own nearest neighbor
50//!     let n_queries = 4;
51//!     let queries = dataset.slice(s![0..n_queries, ..]);
52//!
53//!     let k = 10;
54//!
55//!     // Ivf-Flat search API requires queries and outputs to be on device memory
56//!     // copy query data over, and allocate new device memory for the distances/ neighbors
57//!     // outputs
58//!     let queries = ManagedTensor::from(&queries).to_device(&res)?;
59//!     let mut neighbors_host = ndarray::Array::<u32, _>::zeros((n_queries, k));
60//!     let neighbors = ManagedTensor::from(&neighbors_host).to_device(&res)?;
61//!
62//!     let mut distances_host = ndarray::Array::<f32, _>::zeros((n_queries, k));
63//!     let distances = ManagedTensor::from(&distances_host).to_device(&res)?;
64//!
65//!     let search_params = SearchParams::new()?;
66//!
67//!     index.search(&res, &search_params, &queries, &neighbors, &distances)?;
68//!
69//!     // Copy back to host memory
70//!     distances.to_host(&res, &mut distances_host)?;
71//!     neighbors.to_host(&res, &mut neighbors_host)?;
72//!
73//!     // nearest neighbors should be themselves, since queries are from the
74//!     // dataset
75//!     println!("Neighbors {:?}", neighbors_host);
76//!     println!("Distances {:?}", distances_host);
77//!     Ok(())
78//! }
79//! ```
80
81mod index;
82mod index_params;
83mod search_params;
84
85pub use index::Index;
86pub use index_params::IndexParams;
87pub use search_params::SearchParams;
hipvs/ivf_flat/mod.rs

hipvs/ivf_flat/
mod.rs
