From 44cac0b1c0fbceccaf6ddf04c0564cd20ab758cb Mon Sep 17 00:00:00 2001
From: Siwen Yu <yusiwen@gmail.com>
Date: Tue, 24 Oct 2023 15:31:14 +0800
Subject: [PATCH] chore: add support for building on jetson device (#620)

* add support for building on jetson devices

* fix wrong file to check

* refactor the environment check of CTranslate2 building

* chore: remove unused code

* chore: format code
---
 crates/ctranslate2-bindings/Cargo.toml |  1 +
 crates/ctranslate2-bindings/build.rs   | 20 ++++++++++++++++----
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/crates/ctranslate2-bindings/Cargo.toml b/crates/ctranslate2-bindings/Cargo.toml
index 7cd7168..88d260b 100644
--- a/crates/ctranslate2-bindings/Cargo.toml
+++ b/crates/ctranslate2-bindings/Cargo.toml
@@ -22,3 +22,4 @@ rust-cxx-cmake-bridge = { path = "../rust-cxx-cmake-bridge", optional = true }
 [features]
 default = ["dep:cmake", "dep:rust-cxx-cmake-bridge"]
 link_shared = []
+link_static_cuda = []
diff --git a/crates/ctranslate2-bindings/build.rs b/crates/ctranslate2-bindings/build.rs
index 96c69d9..b805e10 100644
--- a/crates/ctranslate2-bindings/build.rs
+++ b/crates/ctranslate2-bindings/build.rs
@@ -37,16 +37,28 @@ fn link_static() -> PathBuf {
     if cfg!(target_os = "linux") {
         config
             .define("WITH_MKL", "OFF")
-            .define("OPENMP_RUNTIME", "NONE")
-            .define("WITH_OPENBLAS", "ON")
-            .cxxflag("-msse4.1")
+            .define("OPENMP_RUNTIME", "NONE");
+
+        if cfg!(target_feature = "sse4.1") {
+            config.cxxflag("-msse4.1");
+        }
+
+        if cfg!(feature = "link_static_cuda") {
+            config.define("WITH_CUDA", "ON").define("WITH_CUDNN", "ON");
+
+            if cfg!(target_arch = "aarch64") {
+                config.cxxflag("-mcpu=native");
+            }
+        } else {
+            config.define("WITH_OPENBLAS", "ON");
+        }
     } else if cfg!(target_os = "macos") {
         config
             .define("CMAKE_OSX_ARCHITECTURES", "arm64")
             .define("WITH_ACCELERATE", "ON")
             .define("WITH_MKL", "OFF")
             .define("OPENMP_RUNTIME", "NONE")
-            .define("WITH_RUY", "ON")
+            .define("WITH_RUY", "ON");
     } else {
         panic!("Invalid target")
     };