From e5407da5268172c5b874294430d1092d6c76aee0 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 release.yml

---
 .github/workflows/release.yml      | 29 ++++++++++++++++++++++++++---
 ci/prepare_build_environment.ps1   |  1 +
 crates/llama-cpp-bindings/build.rs | 18 ++++++++++++------
 3 files changed, 39 insertions(+), 9 deletions(-)
 create mode 100644 ci/prepare_build_environment.ps1

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index b22676562baf..97fa47242991 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -26,7 +26,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
@@ -40,6 +40,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
@@ -77,13 +81,32 @@ 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: Install CUDA toolkit for Windows
+        uses: Jimver/cuda-toolkit@v0.2.11
+        with:
+          cuda: '11.7.1'
+          method: 'network'
+          sub-packages: '["nvcc", "cudart", "cublas", "cublas_dev", "thrust", "visual_studio_integration"]'
+        if: runner.os == 'Windows'
 
       - name: Bulid release binary
         run: cargo build ${{ matrix.build_args }} --release --target ${{ matrix.target }} --package tabby
 
-      - name: Rename release binary
+      - name: Rename release binary for macOS & Linux
         run: mv target/${{ matrix.target }}/release/tabby tabby_${{ matrix.binary }}
+        if: runner.os != 'Windows'
+
+      - name: Rename release binary for Windows
+        run: mv target/${{ matrix.target }}/release/tabby.exe tabby_${{ matrix.binary }}.exe
+        if: runner.os == 'Windows'
 
       - name: Upload artifacts
         uses: actions/upload-artifact@v3
diff --git a/ci/prepare_build_environment.ps1 b/ci/prepare_build_environment.ps1
new file mode 100644
index 000000000000..7235782df0a4
--- /dev/null
+++ b/ci/prepare_build_environment.ps1
@@ -0,0 +1 @@
+choco install --yes protoc
diff --git a/crates/llama-cpp-bindings/build.rs b/crates/llama-cpp-bindings/build.rs
index ec27dfd8cd0b..a768e361faf2 100644
--- a/crates/llama-cpp-bindings/build.rs
+++ b/crates/llama-cpp-bindings/build.rs
@@ -12,6 +12,8 @@ fn main() {
 
     println!("cargo:rerun-if-changed=include/engine.h");
     println!("cargo:rerun-if-changed=src/engine.cc");
+    println!("cargo:rustc-link-lib=llama");
+    println!("cargo:rustc-link-lib=ggml_static");
 
     build_llama_cpp();
     build_cxx_binding();
@@ -29,9 +31,16 @@ 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=culibos");
+        }
         println!("cargo:rustc-link-lib=cudart");
-        println!("cargo:rustc-link-lib=culibos");
         println!("cargo:rustc-link-lib=cublas");
         println!("cargo:rustc-link-lib=cublasLt");
     }
@@ -82,7 +91,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()
         );
@@ -92,9 +101,6 @@ fn build_llama_cpp() {
 }
 
 fn build_cxx_binding() {
-    println!("cargo:rustc-link-lib=llama");
-    println!("cargo:rustc-link-lib=ggml_static");
-
     cxx_build::bridge("src/lib.rs")
         .file("src/engine.cc")
         .flag_if_supported("-Iinclude")