From 6fbbfb5e934d3a5c9ac24b7cce2b661bd885618f Mon Sep 17 00:00:00 2001
From: darknight <illuminating.me@gmail.com>
Date: Sun, 10 Dec 2023 18:09:07 +0800
Subject: [PATCH] update cuda config, add window release setting

---
 .github/workflows/release.yml      | 14 ++++++++++++--
 ci/prepare_build_environment.ps1   |  2 ++
 crates/llama-cpp-bindings/build.rs | 11 +++++++++--
 3 files changed, 23 insertions(+), 4 deletions(-)
 create mode 100644 ci/prepare_build_environment.ps1

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 5e0e5912c8ea..4fda23368f64 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -24,7 +24,7 @@ jobs:
     container: ${{ matrix.container }}
     strategy:
       matrix:
-        binary: [aarch64-apple-darwin, x86_64-manylinux2014, x86_64-manylinux2014-cuda117]
+        binary: [aarch64-apple-darwin, x86_64-manylinux2014, x86_64-manylinux2014-cuda117, x86_64-pc-windows-msvc]
         include:
           - os: macos-latest
             target: aarch64-apple-darwin
@@ -38,6 +38,10 @@ jobs:
             binary: x86_64-manylinux2014-cuda117
             container: sameli/manylinux2014_x86_64_cuda_11.7
             build_args: --features cuda
+          - os: windows-latest
+            target: x86_64-pc-windows-msvc
+            binary: x86_64-pc-windows-msvc
+            build_args: --features cuda
 
     env:
       SCCACHE_GHA_ENABLED: true
@@ -75,7 +79,13 @@ jobs:
             ~/.cargo/registry
             ~/.cargo/git
 
-      - run: bash ./ci/prepare_build_environment.sh
+      - name: Prepare build environment for macOS & Linux
+        run: bash ./ci/prepare_build_environment.sh
+        if: runner.os != 'Windows'
+
+      - name: Prepare build environment for Windows
+        run: ./ci/prepare_build_environment.ps1
+        if: runner.os == 'Windows'
 
       - name: Bulid release binary
         run: cargo build ${{ matrix.build_args }} --release --target ${{ matrix.target }} --package tabby
diff --git a/ci/prepare_build_environment.ps1 b/ci/prepare_build_environment.ps1
new file mode 100644
index 000000000000..d46e4c2f55ad
--- /dev/null
+++ b/ci/prepare_build_environment.ps1
@@ -0,0 +1,2 @@
+Echo "install protocol buffer compiler..."
+choco install protoc
diff --git a/crates/llama-cpp-bindings/build.rs b/crates/llama-cpp-bindings/build.rs
index ec27dfd8cd0b..ae71658e58db 100644
--- a/crates/llama-cpp-bindings/build.rs
+++ b/crates/llama-cpp-bindings/build.rs
@@ -29,7 +29,14 @@ fn build_llama_cpp() {
     if cfg!(feature = "cuda") {
         config.define("LLAMA_CUBLAS", "ON");
         config.define("CMAKE_POSITION_INDEPENDENT_CODE", "ON");
-        println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64");
+        if cfg!(target_os = "windows") {
+            let Ok(cuda_path) = env::var("CUDA_PATH") else {
+                panic!("CUDA_PATH is not set");
+            };
+            println!(r"cargo:rustc-link-search=native={}\lib\x64", cuda_path);
+        } else {
+            println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64");
+        }
         println!("cargo:rustc-link-lib=cudart");
         println!("cargo:rustc-link-lib=culibos");
         println!("cargo:rustc-link-lib=cublas");
@@ -82,7 +89,7 @@ fn build_llama_cpp() {
     let dst = config.build();
     if cfg!(target_os = "windows") {
         println!(
-            "cargo:rustc-link-search=native={}\\build\\{}",
+            r"cargo:rustc-link-search=native={}\build\{}",
             dst.display(),
             config.get_profile()
         );