Rust bindings for cuvs

rapidsai · AyodeAwe · Mar 4, 2024 · Feb 6, 2024 · Feb 6, 2024 · Feb 7, 2024
commit a4ec121d78bc587b27c71d492ef492f10d6001ab
@@ -0,0 +1,6 @@
+[workspace]
+members = [
+    "cuvs",
+    "cuvs-sys",
+]
+resolver = "2"
@@ -0,0 +1,17 @@
+[package]
+name = "cuvs-sys"
+version = "0.1.0"
+edition = "2021"
+repository = "https://github.com/rapidsai/cuvs"
+homepage = "https://github.com/rapidsai/cuvs"
+description = "Low-level rust bindings to libcuvs"
+authors = ["NVIDIA Corporation"]
+license = "Apache-2.0"
+links = "cuvs"
+
+
+[dependencies]
+
+[build-dependencies]
+cmake = ">=0.1"
+bindgen = ">=0.69"
@@ -0,0 +1,107 @@
+/*
+ * 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");
+
+    // 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<String> = 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).*")
+        .allowlist_function("(cuvs|cagra).*")
+        .rustified_enum("(cuvs|cagra|DL).*")
+        .generate()
+        .expect("Unable to generate cagra_c bindings")
+        .write_to_file(out_path.join("cuvs_bindings.rs"))
+        .expect("Failed to write generated rust bindings");
+}
@@ -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 <cuvs/core/c_api.h>
+#include <cuvs/neighbors/cagra_c.h>
@@ -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::<cagraIndex_t>::uninit();
+            assert_eq!(
+                cagraIndexCreate(index.as_mut_ptr()),
+                cuvsError_t::CUVS_SUCCESS
+            );
+            let index = index.assume_init();
+            assert_eq!(cagraIndexDestroy(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);
+        }
+    }
+}
@@ -0,0 +1,12 @@
+[package]
+name = "cuvs"
+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"
+
+[dependencies]
+ffi = { package = "cuvs-sys", path = "../cuvs-sys" }
@@ -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);
+        }
+    }
+}
@@ -0,0 +1,109 @@
+/*
+ * 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::cagraGraphBuildAlgo;
+
+/// Supplemental parameters to build CAGRA Index
+pub struct IndexParams {
+    pub params: ffi::cuvsCagraIndexParams_t,
+}
+
+impl IndexParams {
+    pub fn new() -> Result<IndexParams> {
+        unsafe {
+            let mut params = core::mem::MaybeUninit::<ffi::cuvsCagraIndexParams_t>::uninit();
+            check_cuvs(ffi::cuvsCagraIndexParamsCreate(params.as_mut_ptr()))?;
+            Ok(IndexParams {
+                params: params.assume_init(),
+            })
+        }
+    }
+
+    /// Degree of input graph for pruning
+    pub fn set_intermediate_graph_degree(self, intermediate_graph_degree: usize) -> IndexParams {
+        unsafe {
+            (*self.params).intermediate_graph_degree = intermediate_graph_degree;
+        }
+        self
+    }
+
+    /// Degree of output graph
+    pub fn set_graph_degree(self, graph_degree: usize) -> IndexParams {
+        unsafe {
+            (*self.params).graph_degree = graph_degree;
+        }
+        self
+    }
+
+    /// ANN algorithm to build knn graph
+    pub fn set_build_algo(self, build_algo: BuildAlgo) -> IndexParams {
+        unsafe {
+            (*self.params).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.params).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.params })
+    }
+}
+
+impl Drop for IndexParams {
+    fn drop(&mut self) {
+        if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraIndexParamsDestroy(self.params) }) {
+            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
+        assert_eq!(format!("{:?}", params), "IndexParams { params: cagraIndexParams { intermediate_graph_degree: 128, graph_degree: 16, build_algo: NN_DESCENT, nn_descent_niter: 10 } }");
+    }
+}
@@ -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};