Update build.rs

TabbyML · Dec 13, 2023 · 18df1ba · 18df1ba
1 parent 5c1ea2f
commit 18df1ba
Showing 1 changed file with 5 additions and 7 deletions.
diff --git a/crates/llama-cpp-bindings/build.rs b/crates/llama-cpp-bindings/build.rs
@@ -45,7 +45,7 @@ fn build_llama_cpp() {
         println!("cargo:rustc-link-lib=cublasLt");
     }
     if cfg!(feature = "rocm") {
-        let amd_gpu_default_targets: Vec<&str> = vec![
+        let amd_gpu_targets: Vec<&str> = vec![
             "gfx803",
             "gfx900",
             "gfx906:xnack-",
@@ -63,17 +63,15 @@ fn build_llama_cpp() {
             "gfx1102",
             "gfx1103",
         ];
-        let amd_gpu_targets =
-            env::var("AMDGPU_TARGETS").unwrap_or(amd_gpu_default_targets.join(";"));
-
+
         let rocm_root = env::var("ROCM_ROOT").unwrap_or("/opt/rocm".to_string());
         config.define("LLAMA_HIPBLAS", "ON");
         config.define("CMAKE_C_COMPILER", format!("{}/llvm/bin/clang", rocm_root));
         config.define(
             "CMAKE_CXX_COMPILER",
             format!("{}/llvm/bin/clang++", rocm_root),
         );
-        config.define("AMDGPU_TARGETS", amd_gpu_targets);
+        config.define("AMDGPU_TARGETS", amd_gpu_targets.join(";"));
         println!("cargo:rustc-link-arg=-Wl,--copy-dt-needed-entries");
         println!("cargo:rustc-link-search=native={}/hip/lib", rocm_root);
         println!("cargo:rustc-link-search=native={}/rocblas/lib", rocm_root);
@@ -105,8 +103,8 @@ fn build_llama_cpp() {
 fn build_cxx_binding() {
     cxx_build::bridge("src/lib.rs")
         .file("src/engine.cc")
-        .include("include")
-        .include("llama.cpp")
+        .flag_if_supported("-Iinclude")
+        .flag_if_supported("-Illama.cpp")
         .flag_if_supported("-std=c++14")
         .compile("cxxbridge");
 }