1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
/*
* SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
* SPDX-License-Identifier: Apache-2.0
*/
use std::env;
use std::io::BufRead;
use std::path::PathBuf;
fn main() {
// build the cuvs c-api library with cmake, and link it into this crate
let cuvs_build = cmake::Config::new(".").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<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")
// needed to find cudaruntime.h
.clang_args(cmake_cxx_flags.split(' '))
// include cuvs c headers and dlpack headers we copied
// into our staging location
.clang_arg(format!("-I{}/build/bindings/include/", out_path.display()))
// 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|bruteForce|cagra|DL).*")
.allowlist_function("(cuvs|bruteForce|cagra).*")
.rustified_enum("(cuvs|cagra|DL|DistanceType|codebook_gen|cudaDataType_t).*")
// also need some basic cuda mem functions for copying data
.allowlist_function("(cudaMemcpyAsync|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");
}