Merge pull request #51 from tazz4843/cuda-and-opencl-support

Add CUDA and OpenCL support

Cargo.toml

@@ -20,8 +20,12 @@ whisper-rs-sys = { path = "sys", version = "0.5" }
 hound = "3.5.0"
 
 [features]
+default = []
+
 simd = []
 coreml = ["whisper-rs-sys/coreml"]
+cuda = ["whisper-rs-sys/cuda"]
+opencl = ["whisper-rs-sys/opencl"]
 test-with-tiny-model = []
 
 [package.metadata.docs.rs]

sys/Cargo.toml

@@ -27,8 +27,11 @@ include = [
 
 [features]
 coreml = []
+cuda = []
+opencl = []
 
 [dependencies]
 
 [build-dependencies]
 bindgen = "0.64"
+cfg-if = "1"

sys/build.rs

@@ -2,6 +2,7 @@
 
 extern crate bindgen;
 
+use cfg_if::cfg_if;
 use std::env;
 use std::path::PathBuf;
 
@@ -25,6 +26,20 @@ fn main() {
     println!("cargo:rustc-link-lib=static=whisper");
     #[cfg(feature = "coreml")]
     println!("cargo:rustc-link-lib=static=whisper.coreml");
+    #[cfg(feature = "opencl")]
+    {
+        println!("cargo:rustc-link-lib=clblast");
+        println!("cargo:rustc-link-lib=OpenCL");
+    }
+    #[cfg(feature = "cuda")]
+    {
+        println!("cargo:rustc-link-lib=cublas");
+        println!("cargo:rustc-link-lib=culibos");
+        println!("cargo:rustc-link-lib=cudart");
+        println!("cargo:rustc-link-lib=cublasLt");
+        println!("cargo:rustc-link-search=/usr/local/cuda/lib64");
+        println!("cargo:rustc-link-search=/opt/cuda/lib64");
+    }
     println!("cargo:rerun-if-changed=wrapper.h");
 
     if env::var("WHISPER_DONT_GENERATE_BINDINGS").is_ok() {
@@ -70,33 +85,28 @@ fn main() {
     _ = std::fs::create_dir("build");
     env::set_current_dir("build").expect("Unable to change directory to whisper.cpp build");
 
-    #[cfg(feature = "coreml")]
+    let mut cmd = std::process::Command::new("cmake");
-    let code = std::process::Command::new("cmake")
+    cmd.arg("..")
-        .arg("..")
         .arg("-DCMAKE_BUILD_TYPE=Release")
         .arg("-DBUILD_SHARED_LIBS=OFF")
         .arg("-DWHISPER_ALL_WARNINGS=OFF")
         .arg("-DWHISPER_ALL_WARNINGS_3RD_PARTY=OFF")
         .arg("-DWHISPER_BUILD_TESTS=OFF")
-        .arg("-DWHISPER_BUILD_EXAMPLES=OFF")
+        .arg("-DWHISPER_BUILD_EXAMPLES=OFF");
-        .arg("-DWHISPER_COREML=1")
-        .arg("-DWHISPER_COREML_ALLOW_FALLBACK=1")
-        .status()
-        .expect("Failed to generate build script");
 
-    #[cfg(not(feature = "coreml"))]
+    #[cfg(feature = "coreml")]
-    let code = std::process::Command::new("cmake")
+    cmd.arg("-DWHISPER_COREML=ON")
-        .arg("..")
+        .arg("-DWHISPER_COREML_ALLOW_FALLBACK=1");
-        .arg("-DCMAKE_BUILD_TYPE=Release")
+
-        .arg("-DBUILD_SHARED_LIBS=OFF")
+    #[cfg(feature = "cuda")]
-        .arg("-DWHISPER_ALL_WARNINGS=OFF")
+    cmd.arg("-DWHISPER_CUBLAS=ON");
-        .arg("-DWHISPER_ALL_WARNINGS_3RD_PARTY=OFF")
+
-        .arg("-DWHISPER_BUILD_TESTS=OFF")
+    #[cfg(feature = "opencl")]
-        .arg("-DWHISPER_BUILD_EXAMPLES=OFF")
+    cmd.arg("-DWHISPER_CLBLAST=ON");
-        .status()
+
-        .expect("Failed to generate build script");
+    let code = cmd.status().expect("Failed to run `cmake`");
     if code.code() != Some(0) {
-        panic!("Failed to generate build script");
+        panic!("Failed to run `cmake`");
     }
 
     let code = std::process::Command::new("cmake")
@@ -111,25 +121,25 @@ fn main() {
     }
 
     // move libwhisper.a to where Cargo expects it (OUT_DIR)
-    #[cfg(target_os = "windows")]
+    cfg_if! {
-    {
+        if #[cfg(target_os = "windows")] {
             std::fs::copy(
                 "Release/whisper.lib",
                 format!("{}/whisper.lib", env::var("OUT_DIR").unwrap()),
             )
             .expect("Failed to copy libwhisper.lib");
-    }
+        } else {
-
-    #[cfg(not(target_os = "windows"))]
-    {
             std::fs::copy(
                 "libwhisper.a",
                 format!("{}/libwhisper.a", env::var("OUT_DIR").unwrap()),
             )
             .expect("Failed to copy libwhisper.a");
         }
-    #[cfg(feature = "coreml")]
+    }
-    #[cfg(not(target_os = "windows"))]
+
+    // if on iOS or macOS, with coreml feature enabled, copy libwhisper.coreml.a as well
+    cfg_if! {
+        if #[cfg(all(feature = "coreml", any(target_os = "ios", target_os = "macos")))]
         {
             std::fs::copy(
                 "libwhisper.coreml.a",
@@ -137,6 +147,7 @@ fn main() {
             )
             .expect("Failed to copy libwhisper.coreml.a");
         }
+    }
 
     // clean the whisper build directory to prevent Cargo from complaining during crate publish
     env::set_current_dir("..").expect("Unable to change directory to whisper.cpp");