Add feature cuda for better ergonomics

Author: haixuanTao

README.md

@@ -11,8 +11,8 @@ This is an attempt at a Rust wrapper for
 
 This project consist on two crates:
 
-* [`onnxruntime-sys`](onnxruntime-sys): Low-level binding to the C API;
-* [`onnxruntime`](onnxruntime): High-level and safe API.
+- [`onnxruntime-sys`](onnxruntime-sys): Low-level binding to the C API;
+- [`onnxruntime`](onnxruntime): High-level and safe API.
 
 [Changelog](CHANGELOG.md)
 
@@ -21,25 +21,25 @@ which provides the following targets:
 
 CPU:
 
-* Linux x86_64
-* macOS x86_64
-* macOS aarch64 (no pre-built binaries, no CI testing, see [#74](https://github.com/nbigaouette/onnxruntime-rs/pull/74))
-* Windows i686
-* Windows x86_64
+- Linux x86_64
+- macOS x86_64
+- macOS aarch64 (no pre-built binaries, no CI testing, see [#74](https://github.com/nbigaouette/onnxruntime-rs/pull/74))
+- Windows i686
+- Windows x86_64
 
 GPU:
 
-* Linux x86_64
-* Windows x86_64
+- Linux x86_64
+- Windows x86_64
 
 ---
 
 **WARNING**:
 
-* This is an experiment and work in progress; it is _not_ complete/working/safe. Help welcome!
-* Basic inference works, see [`onnxruntime/examples/sample.rs`](onnxruntime/examples/sample.rs) or [`onnxruntime/tests/integration_tests.rs`](onnxruntime/tests/integration_tests.rs)
-* ONNX Runtime has many options to control the inference process but those options are not yet exposed.
-* This was developed and tested on macOS Catalina. Other platforms should work but have not been tested.
+- This is an experiment and work in progress; it is _not_ complete/working/safe. Help welcome!
+- Basic inference works, see [`onnxruntime/examples/sample.rs`](onnxruntime/examples/sample.rs) or [`onnxruntime/tests/integration_tests.rs`](onnxruntime/tests/integration_tests.rs)
+- ONNX Runtime has many options to control the inference process but those options are not yet exposed.
+- This was developed and tested on macOS Catalina. Other platforms should work but have not been tested.
 
 ---
 
@@ -58,14 +58,14 @@ To select which strategy to use, set the `ORT_STRATEGY` environment variable to:
 3. `compile`: To compile the library
 
 The `download` strategy supports downloading a version of ONNX that supports CUDA. To use this, set the
-environment variable `ORT_USE_CUDA=1` (only supports Linux or Windows).
+feature `cuda` in `Cargo.toml`.
 
 Until the build script allow compilation of the runtime, see the [compilation notes](ONNX_Compilation_Notes.md)
 for some details on the process.
 
 ### Note on using CUDA
 
-To use CUDA you will need to set `ORT_USE_CUDA=1` but also to set your session with the method `use_cuda` as such:
+To use CUDA you will need to set the feature `cuda` but also to set your session with the method `use_cuda` as such:
 
 ```
     let mut session = environment
@@ -86,9 +86,9 @@ dyld: Library not loaded: @rpath/libonnxruntime.1.7.1.dylib
 
 To fix, one can either:
 
-* Set the `LD_LIBRARY_PATH` environment variable to point to the path where the library can be found.
-* Adapt the `.cargo/config` file to contain a linker flag to provide the **full** path:
-  
+- Set the `LD_LIBRARY_PATH` environment variable to point to the path where the library can be found.
+- Adapt the `.cargo/config` file to contain a linker flag to provide the **full** path:
+
   ```toml
   [target.aarch64-apple-darwin]
   rustflags = ["-C", "link-args=-Wl,-rpath,/full/path/to/onnxruntime/lib"]
@@ -269,9 +269,9 @@ instead of the Rust moderation team.
 
 This project is licensed under either of
 
-* Apache License, Version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or
+- Apache License, Version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or
   http://www.apache.org/licenses/LICENSE-2.0)
-* MIT license ([LICENSE-MIT](LICENSE-MIT) or
+- MIT license ([LICENSE-MIT](LICENSE-MIT) or
   http://opensource.org/licenses/MIT)
 
 at your option.

onnxruntime-sys/Cargo.toml

@@ -34,6 +34,7 @@ default = []
 disable-sys-build-script = []
 # Use bindgen to generate bindings in build.rs
 generate-bindings = ["bindgen"]
+cuda = []
 
 [package.metadata.docs.rs]
 # Disable the build.rs on https://docs.rs since it can cause

onnxruntime-sys/build.rs

@@ -28,8 +28,6 @@ const ORT_ENV_STRATEGY: &str = "ORT_STRATEGY";
 /// Name of environment variable that, if present, contains the location of a pre-built library.
 /// Only used if `ORT_STRATEGY=system`.
 const ORT_ENV_SYSTEM_LIB_LOCATION: &str = "ORT_LIB_LOCATION";
-/// Name of environment variable that, if present, controls wether to use CUDA or not.
-const ORT_ENV_GPU: &str = "ORT_USE_CUDA";
 
 /// Subdirectory (of the 'target' directory) into which to extract the prebuilt library.
 const ORT_PREBUILT_EXTRACT_DIR: &str = "onnxruntime";
@@ -54,7 +52,6 @@ fn main() {
     println!("cargo:rustc-link-search=native={}", lib_dir.display());
 
     println!("cargo:rerun-if-env-changed={}", ORT_ENV_STRATEGY);
-    println!("cargo:rerun-if-env-changed={}", ORT_ENV_GPU);
     println!("cargo:rerun-if-env-changed={}", ORT_ENV_SYSTEM_LIB_LOCATION);
 
     generate_bindings(&include_dir);
@@ -280,17 +277,6 @@ enum Accelerator {
     Gpu,
 }
 
-impl FromStr for Accelerator {
-    type Err = String;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s.to_lowercase().as_str() {
-            "1" | "yes" | "true" | "on" => Ok(Accelerator::Gpu),
-            _ => Ok(Accelerator::None),
-        }
-    }
-}
-
 impl OnnxPrebuiltArchive for Accelerator {
     fn as_onnx_str(&self) -> Cow<str> {
         match self {
@@ -353,6 +339,12 @@ impl OnnxPrebuiltArchive for Triplet {
 }
 
 fn prebuilt_archive_url() -> (PathBuf, String) {
+    let accelerator = if cfg!(feature = "cuda") {
+        Accelerator::Gpu
+    } else {
+        Accelerator::None
+    };
+
     let triplet = Triplet {
         os: env::var("CARGO_CFG_TARGET_OS")
             .expect("Unable to get TARGET_OS")
@@ -362,7 +354,7 @@ fn prebuilt_archive_url() -> (PathBuf, String) {
             .expect("Unable to get TARGET_ARCH")
             .parse()
             .unwrap(),
-        accelerator: env::var(ORT_ENV_GPU).unwrap_or_default().parse().unwrap(),
+        accelerator,
     };
 
     let prebuilt_archive = format!(

onnxruntime/Cargo.toml

@@ -41,6 +41,7 @@ model-fetching = ["ureq"]
 # Disable build script; used for https://docs.rs
 disable-sys-build-script = ["onnxruntime-sys/disable-sys-build-script"]
 generate-bindings = ["onnxruntime-sys/generate-bindings"]
+cuda = ["onnxruntime-sys/cuda"]
 
 [package.metadata.docs.rs]
 features = ["disable-sys-build-script", "model-fetching"]

onnxruntime/src/session.rs

@@ -130,14 +130,15 @@ impl<'a> SessionBuilder<'a> {
         Ok(self)
     }
 
-    /// Set the session to use cuda
+    /// Set the session to use cpu
     pub fn use_cpu(self, use_arena: i32) -> Result<SessionBuilder<'a>> {
         unsafe {
             sys::OrtSessionOptionsAppendExecutionProvider_CPU(self.session_options_ptr, use_arena);
         }
         Ok(self)
     }
 
+    #[cfg(feature = "cuda")]
     /// Set the session to use cuda
     pub fn use_cuda(self, device_id: i32) -> Result<SessionBuilder<'a>> {
         unsafe {