Skip to content

feat(rust)!: compile-time dataset lifetime safety for all index types #1869

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zbennett10 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(rust)!: compile-time dataset lifetime safety for all index types #1869

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
183 changes: 108 additions & 75 deletions rust/cuvs/src/brute_force.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,55 +5,83 @@
//! Brute Force KNN

use std::io::{stderr, Write};
use std::marker::PhantomData;

use crate::distance_type::DistanceType;
use crate::dlpack::ManagedTensor;
use crate::error::{check_cuvs, Result};
use crate::resources::Resources;

/// Brute Force KNN Index
/// Brute Force KNN Index.
///
/// The brute force C++ implementation always stores a non-owning view of the
/// dataset, so the index cannot outlive the `ManagedTensor` it was built
/// from. The lifetime parameter `'a` ties the index to that tensor, so the
/// borrow checker catches use-after-free at compile time.
///
/// # Example
///
/// ```no_run
/// # use cuvs::{ManagedTensor, Resources};
/// # use cuvs::brute_force::Index;
/// # use cuvs::distance_type::DistanceType;
/// let res = Resources::new().unwrap();
/// let arr = ndarray::Array::<f32, _>::zeros((64, 8));
/// let tensor = ManagedTensor::from(&arr).to_device(&res).unwrap();
/// let index = Index::build(&res, DistanceType::L2Expanded, None, &tensor).unwrap();
/// // `tensor` must outlive `index`.
/// ```
#[derive(Debug)]
pub struct Index(ffi::cuvsBruteForceIndex_t);
pub struct Index<'a> {
inner: ffi::cuvsBruteForceIndex_t,
_dataset: PhantomData<&'a ()>,
}

impl Index {
/// Builds a new Brute Force KNN Index from the dataset for efficient search.
impl<'a> Index<'a> {
/// Creates a new empty FFI handle.
fn create_handle() -> Result<ffi::cuvsBruteForceIndex_t> {
unsafe {
let mut index = std::mem::MaybeUninit::<ffi::cuvsBruteForceIndex_t>::uninit();
check_cuvs(ffi::cuvsBruteForceIndexCreate(index.as_mut_ptr()))?;
Ok(index.assume_init())
}
}

/// Builds a new Brute Force KNN Index from `dataset`.
///
/// The compiler enforces that `dataset` outlives the returned index,
/// because the C++ brute force implementation stores a non-owning view
/// of the input and would otherwise dangle.
///
/// # Arguments
///
/// * `res` - Resources to use
/// * `metric` - DistanceType to use for building the index
/// * `metric_arg` - Optional value of `p` for Minkowski distances
/// * `dataset` - A row-major matrix on either the host or device to index
pub fn build<T: Into<ManagedTensor>>(
pub fn build(
res: &Resources,
metric: DistanceType,
metric_arg: Option<f32>,
dataset: T,
) -> Result<Index> {
let dataset: ManagedTensor = dataset.into();
let index = Index::new()?;
dataset: &'a ManagedTensor,
) -> Result<Index<'a>> {
let inner = Self::create_handle()?;
unsafe {
check_cuvs(ffi::cuvsBruteForceBuild(
res.0,
dataset.as_ptr(),
metric,
metric_arg.unwrap_or(2.0),
index.0,
inner,
))?;
}
Ok(index)
Ok(Index {
inner,
_dataset: PhantomData,
})
}

/// Creates a new empty index
pub fn new() -> Result<Index> {
unsafe {
let mut index = std::mem::MaybeUninit::<ffi::cuvsBruteForceIndex_t>::uninit();
check_cuvs(ffi::cuvsBruteForceIndexCreate(index.as_mut_ptr()))?;
Ok(Index(index.assume_init()))
}
}

/// Perform a Nearest Neighbors search on the Index
/// Perform a Nearest Neighbors search on the Index.
///
/// # Arguments
///
Expand All @@ -76,7 +104,7 @@ impl Index {

check_cuvs(ffi::cuvsBruteForceSearch(
res.0,
self.0,
self.inner,
queries.as_ptr(),
neighbors.as_ptr(),
distances.as_ptr(),
Expand All @@ -86,9 +114,9 @@ impl Index {
}
}

impl Drop for Index {
impl Drop for Index<'_> {
fn drop(&mut self) {
if let Err(e) = check_cuvs(unsafe { ffi::cuvsBruteForceIndexDestroy(self.0) }) {
if let Err(e) = check_cuvs(unsafe { ffi::cuvsBruteForceIndexDestroy(self.inner) }) {
write!(stderr(), "failed to call bruteForceIndexDestroy {:?}", e)
.expect("failed to write to stderr");
}
Expand All @@ -99,66 +127,60 @@ impl Drop for Index {
mod tests {
use super::*;
use mark_flaky_tests::flaky;
use ndarray::s;
use ndarray::{s, Array2};
use ndarray_rand::rand_distr::Uniform;
use ndarray_rand::RandomExt;

fn test_bfknn(metric: DistanceType) {
let res = Resources::new().unwrap();

// Create a new random dataset to index
let n_datapoints = 16;
let n_features = 8;
let dataset_host =
ndarray::Array::<f32, _>::random((n_datapoints, n_features), Uniform::new(0., 1.0));

let dataset = ManagedTensor::from(&dataset_host).to_device(&res).unwrap();

println!("dataset {:#?}", dataset_host);

// build the brute force index
let index =
Index::build(&res, metric, None, dataset).expect("failed to create brute force index");

res.sync_stream().unwrap();
/// Build a small random host dataset for the brute-force tests.
fn make_dataset(n_datapoints: usize, n_features: usize) -> Array2<f32> {
ndarray::Array::<f32, _>::random((n_datapoints, n_features), Uniform::new(0., 1.0))
}

// use the first 4 points from the dataset as queries : will test that we get them back
// as their own nearest neighbor
let n_queries = 4;
/// Search `index` for the first `n_queries` rows of `dataset_host` and
/// assert each query is its own nearest neighbor.
fn search_and_verify_self_neighbors(
res: &Resources,
index: &Index<'_>,
dataset_host: &Array2<f32>,
n_queries: usize,
k: usize,
) {
let queries = dataset_host.slice(s![0..n_queries, ..]);
let queries = ManagedTensor::from(&queries).to_device(res).unwrap();

let k = 4;

println!("queries! {:#?}", queries);
let queries = ManagedTensor::from(&queries).to_device(&res).unwrap();
let mut neighbors_host = ndarray::Array::<i64, _>::zeros((n_queries, k));
let neighbors = ManagedTensor::from(&neighbors_host)
.to_device(&res)
.unwrap();
let neighbors = ManagedTensor::from(&neighbors_host).to_device(res).unwrap();

let mut distances_host = ndarray::Array::<f32, _>::zeros((n_queries, k));
let distances = ManagedTensor::from(&distances_host)
.to_device(&res)
.unwrap();
let distances = ManagedTensor::from(&distances_host).to_device(res).unwrap();

index
.search(&res, &queries, &neighbors, &distances)
.unwrap();
.search(res, &queries, &neighbors, &distances)
.expect("search failed");

// Copy back to host memory
distances.to_host(&res, &mut distances_host).unwrap();
neighbors.to_host(&res, &mut neighbors_host).unwrap();
distances.to_host(res, &mut distances_host).unwrap();
neighbors.to_host(res, &mut neighbors_host).unwrap();
res.sync_stream().unwrap();

println!("distances {:#?}", distances_host);
println!("neighbors {:#?}", neighbors_host);
for i in 0..n_queries {
assert_eq!(
neighbors_host[[i, 0]],
i as i64,
"query {i} should be its own nearest neighbor"
);
}
}

// nearest neighbors should be themselves, since queries are from the
// dataset
assert_eq!(neighbors_host[[0, 0]], 0);
assert_eq!(neighbors_host[[1, 0]], 1);
assert_eq!(neighbors_host[[2, 0]], 2);
assert_eq!(neighbors_host[[3, 0]], 3);
fn test_bfknn(metric: DistanceType) {
let res = Resources::new().unwrap();
let dataset_host = make_dataset(16, 8);
let dataset = ManagedTensor::from(&dataset_host).to_device(&res).unwrap();

let index =
Index::build(&res, metric, None, &dataset).expect("failed to build brute force index");
res.sync_stream().unwrap();

search_and_verify_self_neighbors(&res, &index, &dataset_host, 4, 4);
}

/*
Expand All @@ -173,10 +195,21 @@ mod tests {
test_bfknn(DistanceType::L2Expanded);
}

// NOTE: brute_force multiple-search test is omitted here because the C++
// brute_force::index stores a non-owning view into the dataset. Building
// from device data via `build()` drops the ManagedTensor after the call,
// leaving a dangling pointer. A follow-up PR will add dataset lifetime
// enforcement (DatasetOwnership<'a>) to make this safe.
// See: https://github.com/rapidsai/cuvs/issues/1838
/// Validates that a borrowed-dataset brute-force index can be searched
/// multiple times while the dataset is kept alive by the caller.
#[test]
fn test_brute_force_multiple_searches() {
let res = Resources::new().unwrap();
let dataset_host = make_dataset(64, 8);
let dataset_device = ManagedTensor::from(&dataset_host).to_device(&res).unwrap();

let index = Index::build(&res, DistanceType::L2Expanded, None, &dataset_device)
.expect("failed to build brute force index");
res.sync_stream().unwrap();

for _ in 0..3 {
search_and_verify_self_neighbors(&res, &index, &dataset_host, 4, 4);
}
// `dataset_device` is still alive here — the borrow checker enforces it.
}
}
Loading
Loading