use std::env;
use std::path::PathBuf;
fn main() {
let nixl_root_path =
env::var("NIXL_PREFIX").unwrap_or_else(|_| "/opt/nvidia/nvda_nixl".to_string());
let nixl_include_path = format!("{}/include", nixl_root_path);
let arch = match env::var("CARGO_CFG_TARGET_ARCH").unwrap_or_else(|_| "x86_64".to_string()).as_str() {
"x86_64" => "x86_64",
"aarch64" => "aarch64",
other => panic!("Unsupported architecture: {}", other),
};
let nixl_lib_path = format!("{}/lib/{}-linux-gnu", nixl_root_path, arch);
println!("cargo:rustc-link-search={}", nixl_lib_path);
cc::Build::new()
.cpp(true)
.compiler("g++") .file("wrapper.cpp")
.flag("-std=c++17")
.flag("-fPIC")
.include(&nixl_include_path)
.flag("-Wno-unused-parameter")
.flag("-Wno-unused-variable")
.compile("wrapper");
println!("cargo:rustc-link-lib=dylib=nixl");
println!("cargo:rustc-link-lib=dylib=nixl_build");
println!("cargo:rustc-link-lib=dylib=serdes");
println!("cargo:rustc-link-lib=dylib=stdc++");
println!("cargo:rerun-if-changed=wrapper.h");
println!("cargo:rerun-if-changed=wrapper.cpp");
let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
bindgen::Builder::default()
.header("wrapper.h")
.generate()
.expect("Unable to generate bindings")
.write_to_file(out_path.join("bindings.rs"))
.expect("Couldn't write bindings!");
}