From 8e6979f5548e447e155af6d728dc7c29a6e5af79 Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Mon, 4 Mar 2024 11:30:10 -0800 Subject: [PATCH] Add Rust bindings for CAGRA (#34) --- .github/workflows/build.yaml | 13 ++ .github/workflows/pr.yaml | 11 ++ README.md | 61 +++++++++ ci/build_rust.sh | 40 ++++++ dependencies.yaml | 16 ++- rust/Cargo.toml | 16 +++ rust/cuvs-sys/Cargo.toml | 16 +++ rust/cuvs-sys/build.rs | 112 ++++++++++++++++ rust/cuvs-sys/cuvs_c_wrapper.h | 20 +++ rust/cuvs-sys/src/lib.rs | 54 ++++++++ rust/cuvs/Cargo.toml | 16 +++ rust/cuvs/build.rs | 29 ++++ rust/cuvs/examples/cagra.rs | 78 +++++++++++ rust/cuvs/src/cagra/index.rs | 141 ++++++++++++++++++++ rust/cuvs/src/cagra/index_params.rs | 110 ++++++++++++++++ rust/cuvs/src/cagra/mod.rs | 23 ++++ rust/cuvs/src/cagra/search_params.rs | 168 +++++++++++++++++++++++ rust/cuvs/src/dlpack.rs | 190 +++++++++++++++++++++++++++ rust/cuvs/src/error.rs | 51 +++++++ rust/cuvs/src/lib.rs | 24 ++++ rust/cuvs/src/resources.rs | 52 ++++++++ 21 files changed, 1240 insertions(+), 1 deletion(-) create mode 100755 ci/build_rust.sh create mode 100644 rust/Cargo.toml create mode 100644 rust/cuvs-sys/Cargo.toml create mode 100644 rust/cuvs-sys/build.rs create mode 100644 rust/cuvs-sys/cuvs_c_wrapper.h create mode 100644 rust/cuvs-sys/src/lib.rs create mode 100644 rust/cuvs/Cargo.toml create mode 100644 rust/cuvs/build.rs create mode 100644 rust/cuvs/examples/cagra.rs create mode 100644 rust/cuvs/src/cagra/index.rs create mode 100644 rust/cuvs/src/cagra/index_params.rs create mode 100644 rust/cuvs/src/cagra/mod.rs create mode 100644 rust/cuvs/src/cagra/search_params.rs create mode 100644 rust/cuvs/src/dlpack.rs create mode 100644 rust/cuvs/src/error.rs create mode 100644 rust/cuvs/src/lib.rs create mode 100644 rust/cuvs/src/resources.rs diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 34cf1f5b0..61ebdf0bc 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -34,6 +34,19 @@ jobs: branch: ${{ inputs.branch }} date: ${{ inputs.date }} sha: ${{ inputs.sha }} + rust-build: + needs: cpp-build + secrets: inherit + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-24.02 + with: + build_type: ${{ inputs.build_type || 'branch' }} + branch: ${{ inputs.branch }} + arch: "amd64" + date: ${{ inputs.date }} + container_image: "rapidsai/ci-conda:latest" + node_type: "gpu-v100-latest-1" + run_script: "ci/build_rust.sh" + sha: ${{ inputs.sha }} python-build: needs: [cpp-build] secrets: inherit diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 5799f5108..7c6db2c60 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -19,6 +19,7 @@ jobs: - conda-python-build - conda-python-tests - docs-build + - rust-build - wheel-build-cuvs - wheel-tests-cuvs - devcontainer @@ -72,6 +73,16 @@ jobs: arch: "amd64" container_image: "rapidsai/ci-conda:latest" run_script: "ci/build_docs.sh" + rust-build: + needs: conda-cpp-build + secrets: inherit + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-24.02 + with: + build_type: pull-request + node_type: "gpu-v100-latest-1" + arch: "amd64" + container_image: "rapidsai/ci-conda:latest" + run_script: "ci/build_rust.sh" wheel-build-cuvs: needs: checks secrets: inherit diff --git a/README.md b/README.md index dfba9eb4a..14bb02812 100755 --- a/README.md +++ b/README.md @@ -106,6 +106,67 @@ cuvsCagraIndexParamsDestroy(index_params); cuvsResourcesDestroy(res); ``` +### Rust API + +```rust +use cuvs::cagra::{Index, IndexParams, SearchParams}; +use cuvs::{ManagedTensor, Resources, Result}; + +use ndarray::s; +use ndarray_rand::rand_distr::Uniform; +use ndarray_rand::RandomExt; + +/// Example showing how to index and search data with CAGRA +fn cagra_example() -> Result<()> { + let res = Resources::new()?; + + // Create a new random dataset to index + let n_datapoints = 65536; + let n_features = 512; + let dataset = + ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); + + // build the cagra index + let build_params = IndexParams::new()?; + let index = Index::build(&res, &build_params, &dataset)?; + println!( + "Indexed {}x{} datapoints into cagra index", + n_datapoints, n_features + ); + + // 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; + let queries = dataset.slice(s![0..n_queries, ..]); + + let k = 10; + + // CAGRA search API requires queries and outputs to be on device memory + // copy query data over, and allocate new device memory for the distances/ neighbors + // outputs + let queries = ManagedTensor::from(&queries).to_device(&res)?; + let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); + let neighbors = ManagedTensor::from(&neighbors_host).to_device(&res)?; + + let mut distances_host = ndarray::Array::::zeros((n_queries, k)); + let distances = ManagedTensor::from(&distances_host).to_device(&res)?; + + let search_params = SearchParams::new()?; + + index.search(&res, &search_params, &queries, &neighbors, &distances)?; + + // Copy back to host memory + distances.to_host(&res, &mut distances_host)?; + neighbors.to_host(&res, &mut neighbors_host)?; + + // nearest neighbors should be themselves, since queries are from the + // dataset + println!("Neighbors {:?}", neighbors_host); + println!("Distances {:?}", distances_host); + Ok(()) +} +``` + ## Contributing diff --git a/ci/build_rust.sh b/ci/build_rust.sh new file mode 100755 index 000000000..895dd41e0 --- /dev/null +++ b/ci/build_rust.sh @@ -0,0 +1,40 @@ +#!/bin/bash +# Copyright (c) 2024, NVIDIA CORPORATION. + +set -euo pipefail + +rapids-logger "Create test conda environment" +. /opt/conda/etc/profile.d/conda.sh + +rapids-dependency-file-generator \ + --output conda \ + --file_key rust \ + --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" | tee env.yaml + +rapids-mamba-retry env create --force -f env.yaml -n rust + +# seeing failures on activating the environment here on unbound locals +# apply workaround from https://github.com/conda/conda/issues/8186#issuecomment-532874667 +set +eu +conda activate rust +set -eu + +rapids-print-env + +# we need to set up LIBCLANG_PATH to allow rust bindgen to work, +# grab it from the conda env +export LIBCLANG_PATH=$(dirname $(find /opt/conda -name libclang.so | head -n 1)) +echo "LIBCLANG_PATH=$LIBCLANG_PATH" + +rapids-logger "Downloading artifacts from previous jobs" +CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp) + +# installing libcuvs/libraft will speed up the rust build substantially +rapids-mamba-retry install \ + --channel "${CPP_CHANNEL}" \ + libcuvs \ + libraft + +# build and test the rust bindings +cd rust +cargo test diff --git a/dependencies.yaml b/dependencies.yaml index f17b84dff..d7562ce57 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -44,7 +44,12 @@ files: - cupy - docs - py_version - - test_py_cuvs + rust: + output: none + includes: + - build + - cuda + - rust py_build_py_cuvs: output: pyproject pyproject_dir: python/cuvs @@ -308,6 +313,15 @@ dependencies: - recommonmark - sphinx-copybutton - sphinx-markdown-tables + rust: + common: + - output_types: [conda] + packages: + - make + - rust + # clang/liblclang only needed for bindgen support + - clang + - libclang build_wheels: common: - output_types: [requirements, pyproject] diff --git a/rust/Cargo.toml b/rust/Cargo.toml new file mode 100644 index 000000000..7e9bfe4ec --- /dev/null +++ b/rust/Cargo.toml @@ -0,0 +1,16 @@ +[workspace] +members = [ + "cuvs", + "cuvs-sys", +] +resolver = "2" + +[workspace.package] +version = "0.1.0" +edition = "2021" +repository = "https://github.com/rapidsai/cuvs" +homepage = "https://github.com/rapidsai/cuvs" +description = "RAPIDS vector search library" +authors = ["NVIDIA Corporation"] +license = "Apache-2.0" + diff --git a/rust/cuvs-sys/Cargo.toml b/rust/cuvs-sys/Cargo.toml new file mode 100644 index 000000000..b011e6b37 --- /dev/null +++ b/rust/cuvs-sys/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "cuvs-sys" +description = "Low-level rust bindings to libcuvs" +links = "cuvs" +version.workspace = true +edition.workspace = true +repository.workspace = true +homepage.workspace = true +authors.workspace = true +license.workspace = true + +[dependencies] + +[build-dependencies] +cmake = ">=0.1" +bindgen = ">=0.69" diff --git a/rust/cuvs-sys/build.rs b/rust/cuvs-sys/build.rs new file mode 100644 index 000000000..816a6f259 --- /dev/null +++ b/rust/cuvs-sys/build.rs @@ -0,0 +1,112 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +use std::env; +use std::io::BufRead; +use std::path::PathBuf; + +/* + TODO: + * would be nice to use already built versions of libcuvs_c / libcuvs + if they already existed, but this might not be possible here using cmake-rs + (https://github.com/rust-lang/cmake-rs/issues/111) + * figure out how this works with rust packaging: does the c++ code + need to be in a subdirectory? If so would a symlink work here + should we be using static linking ? +*/ +fn main() { + // build the cuvs c-api library with cmake, and link it into this crate + let cuvs_build = cmake::Config::new("../../cpp") + .configure_arg("-DBUILD_TESTS:BOOL=OFF") + .configure_arg("-DBUILD_C_LIBRARY:BOOL=ON") + .build(); + + println!( + "cargo:rustc-link-search=native={}/lib", + cuvs_build.display() + ); + println!("cargo:rustc-link-lib=dylib=cuvs_c"); + println!("cargo:rustc-link-lib=dylib=cudart"); + + // we need some extra flags both to link against cuvs, and also to run bindgen + // specifically we need to: + // * -I flags to set the include path to pick up cudaruntime.h during bindgen + // * -rpath-link settings to link to libraft/libcuvs.so etc during the link + // Rather than redefine the logic to set all these things, lets pick up the values from + // the cuvs cmake build in its CMakeCache.txt and set from there + let out_path = PathBuf::from(env::var("OUT_DIR").unwrap()); + + let cmake_cache: Vec = std::io::BufReader::new( + std::fs::File::open(format!("{}/build/CMakeCache.txt", out_path.display())) + .expect("Failed to open cuvs CMakeCache.txt"), + ) + .lines() + .map(|x| x.expect("Couldn't parse line from CMakeCache.txt")) + .collect(); + + let cmake_cxx_flags = cmake_cache + .iter() + .find(|x| x.starts_with("CMAKE_CXX_FLAGS:STRING=")) + .expect("failed to find CMAKE_CXX_FLAGS in CMakeCache.txt") + .strip_prefix("CMAKE_CXX_FLAGS:STRING=") + .unwrap(); + + let cmake_linker_flags = cmake_cache + .iter() + .find(|x| x.starts_with("CMAKE_EXE_LINKER_FLAGS:STRING=")) + .expect("failed to find CMAKE_EXE_LINKER_FLAGS in CMakeCache.txt") + .strip_prefix("CMAKE_EXE_LINKER_FLAGS:STRING=") + .unwrap(); + + // need to propagate the rpath-link settings to dependent crates =( + // (this will get added as DEP_CUVS_CMAKE_LINKER_ARGS in dependent crates) + println!("cargo:cmake_linker_flags={}", cmake_linker_flags); + + // add the required rpath-link flags to the cargo build + for flag in cmake_linker_flags.split(' ') { + if flag.starts_with("-Wl,-rpath-link") { + println!("cargo:rustc-link-arg={}", flag); + } + } + + // run bindgen to automatically create rust bindings for the cuvs c-api + bindgen::Builder::default() + .header("cuvs_c_wrapper.h") + .clang_arg("-I../../cpp/include") + // needed to find cudaruntime.h + .clang_args(cmake_cxx_flags.split(' ')) + // include dlpack from the cmake build dependencies + .clang_arg(format!( + "-I{}/build/_deps/dlpack-src/include/", + out_path.display() + )) + // add `must_use' declarations to functions returning cuvsError_t + // (so that if you don't check the error code a compile warning is + // generated) + .must_use_type("cuvsError_t") + // Only generate bindings for cuvs/cagra types and functions + .allowlist_type("(cuvs|cagra|DL).*") + .allowlist_function("(cuvs|cagra).*") + .rustified_enum("(cuvs|cagra|DL).*") + // also need some basic cuda mem functions + // (TODO: should we be adding in RMM support instead here?) + .allowlist_function("(cudaMalloc|cudaFree|cudaMemcpy)") + .rustified_enum("cudaError") + .generate() + .expect("Unable to generate cagra_c bindings") + .write_to_file(out_path.join("cuvs_bindings.rs")) + .expect("Failed to write generated rust bindings"); +} diff --git a/rust/cuvs-sys/cuvs_c_wrapper.h b/rust/cuvs-sys/cuvs_c_wrapper.h new file mode 100644 index 000000000..ccca82632 --- /dev/null +++ b/rust/cuvs-sys/cuvs_c_wrapper.h @@ -0,0 +1,20 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// wrapper file containing all the C-API's we should automatically be creating rust +// bindings for +#include +#include diff --git a/rust/cuvs-sys/src/lib.rs b/rust/cuvs-sys/src/lib.rs new file mode 100644 index 000000000..8a261a052 --- /dev/null +++ b/rust/cuvs-sys/src/lib.rs @@ -0,0 +1,54 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// ignore warnings from bindgen +#![allow(non_upper_case_globals)] +#![allow(non_camel_case_types)] +#![allow(non_snake_case)] +#![allow(unused_attributes)] + +// include the generated cuvs_bindings.rs file directly in here +// (this file is automatically generated by bindgen in build.rs) +include!(concat!(env!("OUT_DIR"), "/cuvs_bindings.rs")); + +#[cfg(test)] +mod tests { + use super::*; + // some super basic tests here to make sure we can call into the cuvs library + // the actual logic will be tested out through the higher level bindings + + #[test] + fn test_create_cagra_index() { + unsafe { + let mut index = core::mem::MaybeUninit::::uninit(); + assert_eq!( + cuvsCagraIndexCreate(index.as_mut_ptr()), + cuvsError_t::CUVS_SUCCESS + ); + let index = index.assume_init(); + assert_eq!(cuvsCagraIndexDestroy(index), cuvsError_t::CUVS_SUCCESS); + } + } + + #[test] + fn test_create_resources() { + unsafe { + let mut res: cuvsResources_t = 0; + assert_eq!(cuvsResourcesCreate(&mut res), cuvsError_t::CUVS_SUCCESS); + assert_eq!(cuvsResourcesDestroy(res), cuvsError_t::CUVS_SUCCESS); + } + } +} diff --git a/rust/cuvs/Cargo.toml b/rust/cuvs/Cargo.toml new file mode 100644 index 000000000..cc52db026 --- /dev/null +++ b/rust/cuvs/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "cuvs" +description = "RAPIDS vector search library" +version.workspace = true +edition.workspace = true +repository.workspace = true +homepage.workspace = true +authors.workspace = true +license.workspace = true + +[dependencies] +ffi = { package = "cuvs-sys", path = "../cuvs-sys" } +ndarray = "0.15" + +[dev-dependencies] +ndarray-rand = "*" diff --git a/rust/cuvs/build.rs b/rust/cuvs/build.rs new file mode 100644 index 000000000..fde400dc4 --- /dev/null +++ b/rust/cuvs/build.rs @@ -0,0 +1,29 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +use std::env; + +fn main() { + // add the required rpath-link flags to the cargo build + // TODO: ... this isn't great, there must be a way to propagate this directly without hacks like + // this + let cmake_linker_flags = env::var("DEP_CUVS_CMAKE_LINKER_FLAGS").unwrap(); + for flag in cmake_linker_flags.split(' ') { + if flag.starts_with("-Wl,-rpath-link") { + println!("cargo:rustc-link-arg={}", flag); + } + } +} diff --git a/rust/cuvs/examples/cagra.rs b/rust/cuvs/examples/cagra.rs new file mode 100644 index 000000000..ccc1466dd --- /dev/null +++ b/rust/cuvs/examples/cagra.rs @@ -0,0 +1,78 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +use cuvs::cagra::{Index, IndexParams, SearchParams}; +use cuvs::{ManagedTensor, Resources, Result}; + +use ndarray::s; +use ndarray_rand::rand_distr::Uniform; +use ndarray_rand::RandomExt; + +/// Example showing how to index and search data with CAGRA +fn cagra_example() -> Result<()> { + let res = Resources::new()?; + + // Create a new random dataset to index + let n_datapoints = 65536; + let n_features = 512; + let dataset = + ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); + + // build the cagra index + let build_params = IndexParams::new()?; + let index = Index::build(&res, &build_params, &dataset)?; + println!( + "Indexed {}x{} datapoints into cagra index", + n_datapoints, n_features + ); + + // 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; + let queries = dataset.slice(s![0..n_queries, ..]); + + let k = 10; + + // CAGRA search API requires queries and outputs to be on device memory + // copy query data over, and allocate new device memory for the distances/ neighbors + // outputs + let queries = ManagedTensor::from(&queries).to_device(&res)?; + let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); + let neighbors = ManagedTensor::from(&neighbors_host).to_device(&res)?; + + let mut distances_host = ndarray::Array::::zeros((n_queries, k)); + let distances = ManagedTensor::from(&distances_host).to_device(&res)?; + + let search_params = SearchParams::new()?; + + index.search(&res, &search_params, &queries, &neighbors, &distances)?; + + // Copy back to host memory + distances.to_host(&res, &mut distances_host)?; + neighbors.to_host(&res, &mut neighbors_host)?; + + // nearest neighbors should be themselves, since queries are from the + // dataset + println!("Neighbors {:?}", neighbors_host); + println!("Distances {:?}", distances_host); + Ok(()) +} + +fn main() { + if let Err(e) = cagra_example() { + println!("Failed to run CAGRA: {:?}", e); + } +} diff --git a/rust/cuvs/src/cagra/index.rs b/rust/cuvs/src/cagra/index.rs new file mode 100644 index 000000000..43f032676 --- /dev/null +++ b/rust/cuvs/src/cagra/index.rs @@ -0,0 +1,141 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +use std::io::{stderr, Write}; + +use crate::cagra::{IndexParams, SearchParams}; +use crate::dlpack::ManagedTensor; +use crate::error::{check_cuvs, Result}; +use crate::resources::Resources; + +#[derive(Debug)] +pub struct Index(ffi::cuvsCagraIndex_t); + +impl Index { + /// Builds a new index + pub fn build>( + res: &Resources, + params: &IndexParams, + dataset: T, + ) -> Result { + let dataset: ManagedTensor = dataset.into(); + let index = Index::new()?; + unsafe { + check_cuvs(ffi::cuvsCagraBuild(res.0, params.0, dataset.as_ptr(), index.0))?; + } + Ok(index) + } + + /// Creates a new empty index + pub fn new() -> Result { + unsafe { + let mut index = core::mem::MaybeUninit::::uninit(); + check_cuvs(ffi::cuvsCagraIndexCreate(index.as_mut_ptr()))?; + Ok(Index(index.assume_init())) + } + } + + pub fn search( + self, + res: &Resources, + params: &SearchParams, + queries: &ManagedTensor, + neighbors: &ManagedTensor, + distances: &ManagedTensor, + ) -> Result<()> { + unsafe { + check_cuvs(ffi::cuvsCagraSearch( + res.0, + params.0, + self.0, + queries.as_ptr(), + neighbors.as_ptr(), + distances.as_ptr(), + )) + } + } +} + +impl Drop for Index { + fn drop(&mut self) { + if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraIndexDestroy(self.0) }) { + write!(stderr(), "failed to call cagraIndexDestroy {:?}", e) + .expect("failed to write to stderr"); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ndarray::s; + use ndarray_rand::rand_distr::Uniform; + use ndarray_rand::RandomExt; + + #[test] + fn test_cagra_index() { + let res = Resources::new().unwrap(); + + // Create a new random dataset to index + let n_datapoints = 256; + let n_features = 16; + let dataset = + ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); + + // build the cagra index + let build_params = IndexParams::new().unwrap(); + let index = + Index::build(&res, &build_params, &dataset).expect("failed to create cagra index"); + + // 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; + let queries = dataset.slice(s![0..n_queries, ..]); + + let k = 10; + + // CAGRA search API requires queries and outputs to be on device memory + // copy query data over, and allocate new device memory for the distances/ neighbors + // outputs + let queries = ManagedTensor::from(&queries).to_device(&res).unwrap(); + let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); + let neighbors = ManagedTensor::from(&neighbors_host) + .to_device(&res) + .unwrap(); + + let mut distances_host = ndarray::Array::::zeros((n_queries, k)); + let distances = ManagedTensor::from(&distances_host) + .to_device(&res) + .unwrap(); + + let search_params = SearchParams::new().unwrap(); + + index + .search(&res, &search_params, &queries, &neighbors, &distances) + .unwrap(); + + // Copy back to host memory + distances.to_host(&res, &mut distances_host).unwrap(); + neighbors.to_host(&res, &mut neighbors_host).unwrap(); + + // 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); + } +} diff --git a/rust/cuvs/src/cagra/index_params.rs b/rust/cuvs/src/cagra/index_params.rs new file mode 100644 index 000000000..656ab4a9c --- /dev/null +++ b/rust/cuvs/src/cagra/index_params.rs @@ -0,0 +1,110 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +use crate::error::{check_cuvs, Result}; +use std::fmt; +use std::io::{stderr, Write}; + +pub type BuildAlgo = ffi::cuvsCagraGraphBuildAlgo; + +/// Supplemental parameters to build CAGRA Index +pub struct IndexParams(pub ffi::cuvsCagraIndexParams_t); + +impl IndexParams { + pub fn new() -> Result { + unsafe { + let mut params = core::mem::MaybeUninit::::uninit(); + check_cuvs(ffi::cuvsCagraIndexParamsCreate(params.as_mut_ptr()))?; + Ok(IndexParams(params.assume_init())) + } + } + + /// Degree of input graph for pruning + pub fn set_intermediate_graph_degree(self, intermediate_graph_degree: usize) -> IndexParams { + unsafe { + (*self.0).intermediate_graph_degree = intermediate_graph_degree; + } + self + } + + /// Degree of output graph + pub fn set_graph_degree(self, graph_degree: usize) -> IndexParams { + unsafe { + (*self.0).graph_degree = graph_degree; + } + self + } + + /// ANN algorithm to build knn graph + pub fn set_build_algo(self, build_algo: BuildAlgo) -> IndexParams { + unsafe { + (*self.0).build_algo = build_algo; + } + self + } + + /// Number of iterations to run if building with NN_DESCENT + pub fn set_nn_descent_niter(self, nn_descent_niter: usize) -> IndexParams { + unsafe { + (*self.0).nn_descent_niter = nn_descent_niter; + } + self + } +} + +impl fmt::Debug for IndexParams { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + // custom debug trait here, default value will show the pointer address + // for the inner params object which isn't that useful. + write!(f, "IndexParams {{ params: {:?} }}", unsafe { *self.0 }) + } +} + +impl Drop for IndexParams { + fn drop(&mut self) { + if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraIndexParamsDestroy(self.0) }) { + write!( + stderr(), + "failed to call cuvsCagraIndexParamsDestroy {:?}", + e + ) + .expect("failed to write to stderr"); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_index_params() { + let params = IndexParams::new() + .unwrap() + .set_intermediate_graph_degree(128) + .set_graph_degree(16) + .set_build_algo(BuildAlgo::NN_DESCENT) + .set_nn_descent_niter(10); + + // make sure the setters actually updated internal representation on the c-struct + unsafe { + assert_eq!((*params.0).graph_degree, 16); + assert_eq!((*params.0).intermediate_graph_degree, 128); + assert_eq!((*params.0).build_algo, BuildAlgo::NN_DESCENT); + assert_eq!((*params.0).nn_descent_niter, 10); + } + } +} diff --git a/rust/cuvs/src/cagra/mod.rs b/rust/cuvs/src/cagra/mod.rs new file mode 100644 index 000000000..55705c27a --- /dev/null +++ b/rust/cuvs/src/cagra/mod.rs @@ -0,0 +1,23 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +mod index; +mod index_params; +mod search_params; + +pub use index::Index; +pub use index_params::{BuildAlgo, IndexParams}; +pub use search_params::{HashMode, SearchAlgo, SearchParams}; diff --git a/rust/cuvs/src/cagra/search_params.rs b/rust/cuvs/src/cagra/search_params.rs new file mode 100644 index 000000000..11ac92bdd --- /dev/null +++ b/rust/cuvs/src/cagra/search_params.rs @@ -0,0 +1,168 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +use crate::error::{check_cuvs, Result}; +use std::fmt; +use std::io::{stderr, Write}; + +pub type SearchAlgo = ffi::cuvsCagraSearchAlgo; +pub type HashMode = ffi::cuvsCagraHashMode; + +/// Supplemental parameters to search CAGRA index +pub struct SearchParams(pub ffi::cuvsCagraSearchParams_t); + +impl SearchParams { + pub fn new() -> Result { + unsafe { + let mut params = core::mem::MaybeUninit::::uninit(); + check_cuvs(ffi::cuvsCagraSearchParamsCreate(params.as_mut_ptr()))?; + Ok(SearchParams(params.assume_init())) + } + } + + /// Maximum number of queries to search at the same time (batch size). Auto select when 0 + pub fn set_max_queries(self, max_queries: usize) -> SearchParams { + unsafe { + (*self.0).max_queries = max_queries; + } + self + } + + /// Number of intermediate search results retained during the search. + /// This is the main knob to adjust trade off between accuracy and search speed. + /// Higher values improve the search accuracy + pub fn set_itopk_size(self, itopk_size: usize) -> SearchParams { + unsafe { + (*self.0).itopk_size = itopk_size; + } + self + } + + /// Upper limit of search iterations. Auto select when 0. + pub fn set_max_iterations(self, max_iterations: usize) -> SearchParams { + unsafe { + (*self.0).max_iterations = max_iterations; + } + self + } + + /// Which search implementation to use. + pub fn set_algo(self, algo: SearchAlgo) -> SearchParams { + unsafe { + (*self.0).algo = algo; + } + self + } + + /// Number of threads used to calculate a single distance. 4, 8, 16, or 32. + pub fn set_team_size(self, team_size: usize) -> SearchParams { + unsafe { + (*self.0).team_size = team_size; + } + self + } + + /// Lower limit of search iterations. + pub fn set_min_iterations(self, min_iterations: usize) -> SearchParams { + unsafe { + (*self.0).min_iterations = min_iterations; + } + self + } + + /// Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. + pub fn set_thread_block_size(self, thread_block_size: usize) -> SearchParams { + unsafe { + (*self.0).thread_block_size = thread_block_size; + } + self + } + + /// Hashmap type. Auto selection when AUTO. + pub fn set_hashmap_mode(self, hashmap_mode: HashMode) -> SearchParams { + unsafe { + (*self.0).hashmap_mode = hashmap_mode; + } + self + } + + /// Lower limit of hashmap bit length. More than 8. + pub fn set_hashmap_min_bitlen(self, hashmap_min_bitlen: usize) -> SearchParams { + unsafe { + (*self.0).hashmap_min_bitlen = hashmap_min_bitlen; + } + self + } + + /// Upper limit of hashmap fill rate. More than 0.1, less than 0.9. + pub fn set_hashmap_max_fill_rate(self, hashmap_max_fill_rate: f32) -> SearchParams { + unsafe { + (*self.0).hashmap_max_fill_rate = hashmap_max_fill_rate; + } + self + } + + /// Number of iterations of initial random seed node selection. 1 or more. + pub fn set_num_random_samplings(self, num_random_samplings: u32) -> SearchParams { + unsafe { + (*self.0).num_random_samplings = num_random_samplings; + } + self + } + + /// Bit mask used for initial random seed node selection. + pub fn set_rand_xor_mask(self, rand_xor_mask: u64) -> SearchParams { + unsafe { + (*self.0).rand_xor_mask = rand_xor_mask; + } + self + } +} + +impl fmt::Debug for SearchParams { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + // custom debug trait here, default value will show the pointer address + // for the inner params object which isn't that useful. + write!(f, "SearchParams {{ params: {:?} }}", unsafe { *self.0 }) + } +} + +impl Drop for SearchParams { + fn drop(&mut self) { + if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraSearchParamsDestroy(self.0) }) { + write!( + stderr(), + "failed to call cuvsCagraSearchParamsDestroy {:?}", + e + ) + .expect("failed to write to stderr"); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_search_params() { + let params = SearchParams::new().unwrap().set_itopk_size(128); + + unsafe { + assert_eq!((*params.0).itopk_size, 128); + } + } +} diff --git a/rust/cuvs/src/dlpack.rs b/rust/cuvs/src/dlpack.rs new file mode 100644 index 000000000..b86959db1 --- /dev/null +++ b/rust/cuvs/src/dlpack.rs @@ -0,0 +1,190 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +use std::convert::From; + +use crate::error::{check_cuda, Result}; +use crate::resources::Resources; + +#[derive(Debug)] +pub struct ManagedTensor(ffi::DLManagedTensor); + +pub trait IntoDtype { + fn ffi_dtype() -> ffi::DLDataType; +} + +impl ManagedTensor { + pub fn as_ptr(&self) -> *mut ffi::DLManagedTensor { + &self.0 as *const _ as *mut _ + } + + fn bytes(&self) -> usize { + // figure out how many bytes to allocate + let mut bytes: usize = 1; + for x in 0..self.0.dl_tensor.ndim { + bytes *= unsafe { (*self.0.dl_tensor.shape.add(x as usize)) as usize }; + } + bytes *= (self.0.dl_tensor.dtype.bits / 8) as usize; + bytes + } + + pub fn to_device(&self, _res: &Resources) -> Result { + unsafe { + let bytes = self.bytes(); + let mut device_data: *mut std::ffi::c_void = std::ptr::null_mut(); + + // allocate storage, copy over + check_cuda(ffi::cudaMalloc(&mut device_data as *mut _, bytes))?; + check_cuda(ffi::cudaMemcpy( + device_data, + self.0.dl_tensor.data, + bytes, + ffi::cudaMemcpyKind_cudaMemcpyDefault, + ))?; + + let mut ret = self.0.clone(); + ret.dl_tensor.data = device_data; + // call cudaFree automatically to clean up data + ret.deleter = Some(cuda_free_tensor); + ret.dl_tensor.device.device_type = ffi::DLDeviceType::kDLCUDA; + + Ok(ManagedTensor(ret)) + } + } + pub fn to_host< + T: IntoDtype, + S: ndarray::RawData + ndarray::RawDataMut, + D: ndarray::Dimension, + >( + &self, + _res: &Resources, + arr: &mut ndarray::ArrayBase, + ) -> Result<()> { + unsafe { + let bytes = self.bytes(); + check_cuda(ffi::cudaMemcpy( + arr.as_mut_ptr() as *mut std::ffi::c_void, + self.0.dl_tensor.data, + bytes, + ffi::cudaMemcpyKind_cudaMemcpyDefault, + ))?; + + Ok(()) + } + } +} + +unsafe extern "C" fn cuda_free_tensor(self_: *mut ffi::DLManagedTensor) { + let _ = ffi::cudaFree((*self_).dl_tensor.data); +} + +/// Create a non-owning view of a Tensor from a ndarray +impl, D: ndarray::Dimension> + From<&ndarray::ArrayBase> for ManagedTensor +{ + fn from(arr: &ndarray::ArrayBase) -> Self { + // There is a draft PR out right now for creating dlpack directly from ndarray + // right now, but until its merged we have to implement ourselves + //https://github.com/rust-ndarray/ndarray/pull/1306/files + unsafe { + let mut ret = std::mem::MaybeUninit::::uninit(); + let tensor = ret.as_mut_ptr(); + (*tensor).data = arr.as_ptr() as *mut std::os::raw::c_void; + (*tensor).device = ffi::DLDevice { + device_type: ffi::DLDeviceType::kDLCPU, + device_id: 0, + }; + (*tensor).byte_offset = 0; + (*tensor).strides = std::ptr::null_mut(); // TODO: error if not rowmajor + (*tensor).ndim = arr.ndim() as i32; + (*tensor).shape = arr.shape().as_ptr() as *mut _; + (*tensor).dtype = T::ffi_dtype(); + ManagedTensor(ffi::DLManagedTensor { + dl_tensor: ret.assume_init(), + manager_ctx: std::ptr::null_mut(), + deleter: None, + }) + } + } +} + +impl Drop for ManagedTensor { + fn drop(&mut self) { + unsafe { + if let Some(deleter) = self.0.deleter { + deleter(&mut self.0 as *mut _); + } + } + } +} + +impl IntoDtype for f32 { + fn ffi_dtype() -> ffi::DLDataType { + ffi::DLDataType { + code: ffi::DLDataTypeCode::kDLFloat as _, + bits: 32, + lanes: 1, + } + } +} + +impl IntoDtype for f64 { + fn ffi_dtype() -> ffi::DLDataType { + ffi::DLDataType { + code: ffi::DLDataTypeCode::kDLFloat as _, + bits: 64, + lanes: 1, + } + } +} + +impl IntoDtype for i32 { + fn ffi_dtype() -> ffi::DLDataType { + ffi::DLDataType { + code: ffi::DLDataTypeCode::kDLInt as _, + bits: 32, + lanes: 1, + } + } +} + +impl IntoDtype for u32 { + fn ffi_dtype() -> ffi::DLDataType { + ffi::DLDataType { + code: ffi::DLDataTypeCode::kDLUInt as _, + bits: 32, + lanes: 1, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_from_ndarray() { + let arr = ndarray::Array::::zeros((8, 4)); + + let tensor = unsafe { (*(ManagedTensor::from(&arr).as_ptr())).dl_tensor }; + + assert_eq!(tensor.ndim, 2); + + // make sure we can get the shape ok + assert_eq!(unsafe { *tensor.shape }, 8); + assert_eq!(unsafe { *tensor.shape.add(1) }, 4); + } +} diff --git a/rust/cuvs/src/error.rs b/rust/cuvs/src/error.rs new file mode 100644 index 000000000..618106aba --- /dev/null +++ b/rust/cuvs/src/error.rs @@ -0,0 +1,51 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +use std::fmt; + +#[derive(Debug, Clone)] +pub enum Error { + CudaError(ffi::cudaError_t), + CuvsError(ffi::cuvsError_t), +} + +impl std::error::Error for Error {} + +pub type Result = std::result::Result; + +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + Error::CudaError(cuda_error) => write!(f, "cudaError={:?}", cuda_error), + Error::CuvsError(cuvs_error) => write!(f, "cuvsError={:?}", cuvs_error), + } + } +} + +/// Simple wrapper to convert a cuvsError_t into a Result +pub fn check_cuvs(err: ffi::cuvsError_t) -> Result<()> { + match err { + ffi::cuvsError_t::CUVS_SUCCESS => Ok(()), + _ => Err(Error::CuvsError(err)), + } +} + +pub fn check_cuda(err: ffi::cudaError_t) -> Result<()> { + match err { + ffi::cudaError::cudaSuccess => Ok(()), + _ => Err(Error::CudaError(err)), + } +} diff --git a/rust/cuvs/src/lib.rs b/rust/cuvs/src/lib.rs new file mode 100644 index 000000000..7a6f847f5 --- /dev/null +++ b/rust/cuvs/src/lib.rs @@ -0,0 +1,24 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +pub mod cagra; +mod dlpack; +mod error; +mod resources; + +pub use dlpack::ManagedTensor; +pub use error::{Error, Result}; +pub use resources::Resources; diff --git a/rust/cuvs/src/resources.rs b/rust/cuvs/src/resources.rs new file mode 100644 index 000000000..ad7113e6b --- /dev/null +++ b/rust/cuvs/src/resources.rs @@ -0,0 +1,52 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +use crate::error::{check_cuvs, Result}; +use std::io::{stderr, Write}; + +#[derive(Debug)] +pub struct Resources(pub ffi::cuvsResources_t); + +impl Resources { + pub fn new() -> Result { + let mut res: ffi::cuvsResources_t = 0; + unsafe { + check_cuvs(ffi::cuvsResourcesCreate(&mut res))?; + } + Ok(Resources(res)) + } +} + +impl Drop for Resources { + fn drop(&mut self) { + unsafe { + if let Err(e) = check_cuvs(ffi::cuvsResourcesDestroy(self.0)) { + write!(stderr(), "failed to call cuvsResourcesDestroy {:?}", e) + .expect("failed to write to stderr"); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_resources_create() { + let _ = Resources::new(); + } +}