Skip to content

Commit

Permalink
Refine CUDA support.
Browse files Browse the repository at this point in the history
Instead of counting all source files we count only .cu files, ones with
device code that actually require the special treatment.
  • Loading branch information
dot-asm authored and thomcc committed Nov 23, 2022
1 parent e0df9ba commit 792c31c
Showing 1 changed file with 9 additions and 2 deletions.
11 changes: 9 additions & 2 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1369,7 +1369,7 @@ impl Build {
if !msvc || !is_asm || !is_arm {
cmd.arg("-c");
}
if self.cuda && self.files.len() > 1 {
if self.cuda && self.cuda_file_count() > 1 {
cmd.arg("--device-c");
}
if is_asm {
Expand Down Expand Up @@ -2037,7 +2037,7 @@ impl Build {
self.assemble_progressive(dst, chunk)?;
}

if self.cuda {
if self.cuda && self.cuda_file_count() > 0 {
// Link the device-side code and add it to the target library,
// so that non-CUDA linker can link the final binary.

Expand Down Expand Up @@ -3012,6 +3012,13 @@ impl Build {
cache.insert(sdk.into(), ret.clone());
Ok(ret)
}

fn cuda_file_count(&self) -> usize {
self.files
.iter()
.filter(|file| file.extension() == Some(OsStr::new("cu")))
.count()
}
}

impl Default for Build {
Expand Down

0 comments on commit 792c31c

Please sign in to comment.