Improve CUDA support (#612)

* Perform CUDA --device-link.

This allows to perform the final link with system linker.

* Add 'cudart' method mimicking the '--cudart' nvcc command-line option.

Try to locate the library in standard location relative to nvcc command.
If it fails, user is held responsible for specifying one in RUSTFLAGS.

* Add dummy CUDA test to cc-test.

Execution is bound to fail without card, but the failure is ignored.
It's rather a compile-n-link test. The test is suppressed if 'nvcc'
is not found on the $PATH.

* Add dummy CUDA CI test.

* Harmonize CUDA support with NVCC default --cudart static.

This can interfere with current deployments in the wild, in which
case some adjustments might be required. Most notably one might
have to add .cuda("none") to the corresponding Builder instantiation
to restore the original behaviour.

Author: dot-asm

.github/workflows/main.yml

@@ -87,6 +87,25 @@ jobs:
     - run: cargo test ${{ matrix.no_run }} --manifest-path cc-test/Cargo.toml --target ${{ matrix.target }} --features parallel
     - run: cargo test ${{ matrix.no_run }} --manifest-path cc-test/Cargo.toml --target ${{ matrix.target }} --release
 
+  cuda:
+    name: Test CUDA support
+    runs-on: ubuntu-20.04
+    steps:
+    - uses: actions/checkout@master
+    - name: Install cuda-minimal-build-11-4
+      shell: bash
+      run: |
+        # https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=20.04&target_type=deb_network
+        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
+        sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
+        sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
+        sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
+        sudo apt-get update
+        sudo apt-get -y install cuda-minimal-build-11-4
+    - name: Test 'cudart' feature
+      shell: bash
+      run: env PATH=/usr/local/cuda/bin:$PATH cargo test --manifest-path cc-test/Cargo.toml --features test_cuda
+
   msrv:
     name: MSRV
     runs-on: ${{ matrix.os }}

README.md

@@ -189,6 +189,7 @@ fn main() {
     cc::Build::new()
         // Switch to CUDA C++ library compilation using NVCC.
         .cuda(true)
+        .cudart("static")
         // Generate code for Maxwell (GTX 970, 980, 980 Ti, Titan X).
         .flag("-gencode").flag("arch=compute_52,code=sm_52")
         // Generate code for Maxwell (Jetson TX1).

cc-test/Cargo.toml

@@ -11,6 +11,8 @@ test = false
 
 [build-dependencies]
 cc = { path = ".." }
+which = "^4.0"
 
 [features]
 parallel = ["cc/parallel"]
+test_cuda = []

cc-test/build.rs

@@ -35,6 +35,24 @@ fn main() {
         .cpp(true)
         .compile("baz");
 
+    if env::var("CARGO_FEATURE_TEST_CUDA").is_ok() {
+        // Detect if there is CUDA compiler and engage "cuda" feature.
+        let nvcc = match env::var("NVCC") {
+            Ok(var) => which::which(var),
+            Err(_) => which::which("nvcc"),
+        };
+        if nvcc.is_ok() {
+            cc::Build::new()
+                .cuda(true)
+                .cudart("static")
+                .file("src/cuda.cu")
+                .compile("libcuda.a");
+
+            // Communicate [cfg(feature = "cuda")] to test/all.rs.
+            println!("cargo:rustc-cfg=feature=\"cuda\"");
+        }
+    }
+
     if target.contains("windows") {
         cc::Build::new().file("src/windows.c").compile("windows");
     }

cc-test/src/cuda.cu

@@ -0,0 +1,5 @@
+#include <cuda.h>
+
+__global__ void kernel() {}
+
+extern "C" void cuda_kernel() { kernel<<<1, 1>>>(); }

cc-test/tests/all.rs

@@ -56,3 +56,14 @@ fn opt_linkage() {
         assert_eq!(answer(), 42);
     }
 }
+
+#[cfg(feature = "cuda")]
+#[test]
+fn cuda_here() {
+    extern "C" {
+        fn cuda_kernel();
+    }
+    unsafe {
+        cuda_kernel();
+    }
+}

src/lib.rs

@@ -103,6 +103,7 @@ pub struct Build {
     cpp_link_stdlib: Option<Option<String>>,
     cpp_set_stdlib: Option<String>,
     cuda: bool,
+    cudart: Option<String>,
     target: Option<String>,
     host: Option<String>,
     out_dir: Option<PathBuf>,
@@ -298,6 +299,7 @@ impl Build {
             cpp_link_stdlib: None,
             cpp_set_stdlib: None,
             cuda: false,
+            cudart: None,
             target: None,
             host: None,
             out_dir: None,
@@ -611,6 +613,20 @@ impl Build {
         self.cuda = cuda;
         if cuda {
             self.cpp = true;
+            self.cudart = Some("static".to_string());
+        }
+        self
+    }
+
+    /// Link CUDA run-time.
+    ///
+    /// This option mimics the `--cudart` NVCC command-line option. Just like
+    /// the original it accepts `{none|shared|static}`, with default being
+    /// `static`. The method has to be invoked after `.cuda(true)`, or not
+    /// at all, if the default is right for the project.
+    pub fn cudart(&mut self, cudart: &str) -> &mut Build {
+        if self.cuda {
+            self.cudart = Some(cudart.to_string());
         }
         self
     }
@@ -996,6 +1012,56 @@ impl Build {
             }
         }
 
+        let cudart = match &self.cudart {
+            Some(opt) => opt.as_str(), // {none|shared|static}
+            None => "none",
+        };
+        if cudart != "none" {
+            if let Some(nvcc) = which(&self.get_compiler().path) {
+                // Try to figure out the -L search path. If it fails,
+                // it's on user to specify one by passing it through
+                // RUSTFLAGS environment variable.
+                let mut libtst = false;
+                let mut libdir = nvcc;
+                libdir.pop(); // remove 'nvcc'
+                libdir.push("..");
+                let target_arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap();
+                if cfg!(target_os = "linux") {
+                    libdir.push("targets");
+                    libdir.push(target_arch.to_owned() + "-linux");
+                    libdir.push("lib");
+                    libtst = true;
+                } else if cfg!(target_env = "msvc") {
+                    libdir.push("lib");
+                    match target_arch.as_str() {
+                        "x86_64" => {
+                            libdir.push("x64");
+                            libtst = true;
+                        }
+                        "x86" => {
+                            libdir.push("Win32");
+                            libtst = true;
+                        }
+                        _ => libtst = false,
+                    }
+                }
+                if libtst && libdir.is_dir() {
+                    println!(
+                        "cargo:rustc-link-search=native={}",
+                        libdir.to_str().unwrap()
+                    );
+                }
+
+                // And now the -l flag.
+                let lib = match cudart {
+                    "shared" => "cudart",
+                    "static" => "cudart_static",
+                    bad => panic!("unsupported cudart option: {}", bad),
+                };
+                println!("cargo:rustc-link-lib={}", lib);
+            }
+        }
+
         Ok(())
     }
 
@@ -1205,6 +1271,9 @@ impl Build {
         if !msvc || !is_asm || !is_arm {
             cmd.arg("-c");
         }
+        if self.cuda && self.files.len() > 1 {
+            cmd.arg("--device-c");
+        }
         cmd.arg(&obj.src);
         if cfg!(target_os = "macos") {
             self.fix_env_for_apple_os(&mut cmd)?;
@@ -1811,6 +1880,21 @@ impl Build {
             self.assemble_progressive(dst, chunk)?;
         }
 
+        if self.cuda {
+            // Link the device-side code and add it to the target library,
+            // so that non-CUDA linker can link the final binary.
+
+            let out_dir = self.get_out_dir()?;
+            let dlink = out_dir.join(lib_name.to_owned() + "_dlink.o");
+            let mut nvcc = self.get_compiler().to_command();
+            nvcc.arg("--device-link")
+                .arg("-o")
+                .arg(dlink.clone())
+                .arg(dst);
+            run(&mut nvcc, "nvcc")?;
+            self.assemble_progressive(dst, &[dlink])?;
+        }
+
         let target = self.get_target()?;
         if target.contains("msvc") {
             // The Rust compiler will look for libfoo.a and foo.lib, but the
@@ -3100,3 +3184,23 @@ fn map_darwin_target_from_rust_to_compiler_architecture(target: &str) -> Option<
         None
     }
 }
+
+fn which(tool: &Path) -> Option<PathBuf> {
+    fn check_exe(exe: &mut PathBuf) -> bool {
+        let exe_ext = std::env::consts::EXE_EXTENSION;
+        exe.exists() || (!exe_ext.is_empty() && exe.set_extension(exe_ext) && exe.exists())
+    }
+
+    // If |tool| is not just one "word," assume it's an actual path...
+    if tool.components().count() > 1 {
+        let mut exe = PathBuf::from(tool);
+        return if check_exe(&mut exe) { Some(exe) } else { None };
+    }
+
+    // Loop through PATH entries searching for the |tool|.
+    let path_entries = env::var_os("PATH")?;
+    env::split_paths(&path_entries).find_map(|path_entry| {
+        let mut exe = path_entry.join(tool);
+        return if check_exe(&mut exe) { Some(exe) } else { None };
+    })
+}