use std::ffi::OsStr;
use std::path::{Path, PathBuf};
use std::process::{Command, Output};
#[inline]
#[must_use]
pub fn nvbit_include() -> PathBuf {
PathBuf::from(std::env::var("DEP_NVBIT_INCLUDE").expect("nvbit include path"))
.canonicalize()
.expect("canonicalize path")
}
#[inline]
#[must_use]
pub fn output_path() -> PathBuf {
PathBuf::from(std::env::var("OUT_DIR").expect("cargo out dir"))
.canonicalize()
.expect("canonicalize path")
}
#[inline]
#[must_use]
pub fn manifest_path() -> PathBuf {
PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").expect("cargo manifest dir"))
.canonicalize()
.expect("canonicalize path")
}
#[derive(thiserror::Error, Debug)]
pub enum Error {
#[error(transparent)]
Io(#[from] std::io::Error),
#[error("Command failed")]
Command(Output),
}
#[derive(Debug, Clone)]
pub struct Build {
include_directories: Vec<PathBuf>,
objects: Vec<PathBuf>,
sources: Vec<PathBuf>,
instrumentation_sources: Vec<PathBuf>,
compiler_flags: Vec<String>,
host_compiler: Option<PathBuf>,
nvcc_compiler: Option<PathBuf>,
warnings: bool,
warnings_as_errors: bool,
}
impl Default for Build {
fn default() -> Self {
Self::new()
}
}
impl Build {
#[inline]
#[must_use]
pub fn new() -> Self {
Self {
include_directories: Vec::new(),
objects: Vec::new(),
sources: Vec::new(),
instrumentation_sources: Vec::new(),
compiler_flags: Vec::new(),
host_compiler: None,
nvcc_compiler: None,
warnings: false,
warnings_as_errors: false,
}
}
fn compile_instrumentation_functions(
&self,
nvcc_compiler: &Path,
include_args: &[String],
compiler_flags: &[&str],
objects: &mut Vec<PathBuf>,
) -> Result<(), Error> {
for (i, src) in self.instrumentation_sources.iter().enumerate() {
let default_name = format!("instr_src_{i}");
let obj = output_path()
.join(
src.file_name()
.and_then(OsStr::to_str)
.unwrap_or(&default_name),
)
.with_extension("o");
let mut cmd = Command::new(nvcc_compiler);
if let Some(host_compiler) = &self.host_compiler {
cmd.args(["-ccbin", &*host_compiler.to_string_lossy()]);
}
cmd.args(include_args);
cmd.args([
"-maxrregcount=24",
"-arch=sm_35",
"-Xptxas",
"-astoolspatch",
"--keep-device-functions",
])
.args(compiler_flags)
.args(&self.compiler_flags)
.arg("-c")
.arg(src)
.arg("-o")
.arg(&*obj.to_string_lossy());
println!("cargo:warning={cmd:?}");
let result = cmd.output()?;
if !result.status.success() {
return Err(Error::Command(result));
}
objects.push(obj);
}
Ok(())
}
fn compile_sources(
&self,
nvcc_compiler: &Path,
include_args: &[String],
compiler_flags: &[&str],
objects: &mut Vec<PathBuf>,
) -> Result<(), Error> {
for (i, src) in self.sources.iter().enumerate() {
let default_name = format!("src_{i}");
let obj = output_path()
.join(
src.file_name()
.and_then(OsStr::to_str)
.unwrap_or(&default_name),
)
.with_extension("o");
let mut cmd = Command::new(nvcc_compiler);
if let Some(host_compiler) = &self.host_compiler {
cmd.args(["-ccbin", &*host_compiler.to_string_lossy()]);
}
cmd.args(include_args)
.args(compiler_flags)
.args(&self.compiler_flags)
.args(["-dc", "-c"])
.arg(src)
.arg("-o")
.arg(&*obj.to_string_lossy());
println!("cargo:warning={cmd:?}");
let result = cmd.output()?;
if !result.status.success() {
return Err(Error::Command(result));
}
objects.push(obj);
}
Ok(())
}
pub fn compile<O: AsRef<str>>(&self, output: O) -> Result<(), Error> {
let mut objects = self.objects.clone();
let include_args: Vec<_> = self
.include_directories
.iter()
.map(|d| format!("-I{}", &d.to_string_lossy()))
.collect();
let mut compiler_flags = vec!["-arch=sm_35", "-Xcompiler", "-fPIC"];
if self.warnings {
compiler_flags.extend(["-Xcompiler", "-Wall"]);
}
if self.warnings_as_errors {
compiler_flags.extend(["-Xcompiler", "-Werror"]);
}
let default_nvcc_compiler = PathBuf::from("nvcc");
let nvcc_compiler = self
.nvcc_compiler
.as_ref()
.unwrap_or(&default_nvcc_compiler);
self.compile_instrumentation_functions(
nvcc_compiler,
&include_args,
&compiler_flags,
&mut objects,
)?;
self.compile_sources(nvcc_compiler, &include_args, &compiler_flags, &mut objects)?;
let dev_link_obj = output_path().join("dev_link.o");
let mut cmd = Command::new(nvcc_compiler);
if let Some(host_compiler) = &self.host_compiler {
cmd.args(["-ccbin", &*host_compiler.to_string_lossy()]);
}
cmd.args(&include_args)
.args(&compiler_flags)
.args(&self.compiler_flags)
.arg("-dlink")
.args(&objects)
.arg("-o")
.arg(&*dev_link_obj.to_string_lossy());
println!("cargo:warning={cmd:?}");
let result = cmd.output()?;
if !result.status.success() {
return Err(Error::Command(result));
}
objects.push(dev_link_obj);
let mut cmd = Command::new("ar");
cmd.args([
"cru",
&output_path()
.join(format!("lib{}.a", output.as_ref()))
.to_string_lossy(),
])
.args(&objects);
println!("cargo:warning={cmd:?}");
let result = cmd.output()?;
if !result.status.success() {
return Err(Error::Command(result));
}
println!("cargo:rustc-link-search=native={}", output_path().display());
println!(
"cargo:rustc-link-lib=static:+whole-archive={}",
output.as_ref()
);
Ok(())
}
pub fn host_compiler<P: Into<PathBuf>>(&mut self, compiler: P) -> &mut Self {
self.host_compiler = Some(compiler.into());
self
}
pub fn nvcc_compiler<P: Into<PathBuf>>(&mut self, compiler: P) -> &mut Self {
self.nvcc_compiler = Some(compiler.into());
self
}
pub fn object<P: Into<PathBuf>>(&mut self, obj: P) -> &mut Self {
self.objects.push(obj.into());
self
}
pub fn objects<P>(&mut self, objects: P) -> &mut Self
where
P: IntoIterator,
P::Item: Into<PathBuf>,
{
for obj in objects {
self.object(obj);
}
self
}
pub fn instrumentation_source<P: Into<PathBuf>>(&mut self, src: P) -> &mut Self {
self.instrumentation_sources.push(src.into());
self
}
pub fn instrumentation_sources<P>(&mut self, sources: P) -> &mut Self
where
P: IntoIterator,
P::Item: Into<PathBuf>,
{
for src in sources {
self.instrumentation_source(src);
}
self
}
pub fn source<P: Into<PathBuf>>(&mut self, dir: P) -> &mut Self {
self.sources.push(dir.into());
self
}
pub fn sources<P>(&mut self, sources: P) -> &mut Self
where
P: IntoIterator,
P::Item: Into<PathBuf>,
{
for src in sources {
self.source(src);
}
self
}
pub fn nvcc_flag<F: Into<String>>(&mut self, flag: F) -> &mut Build {
self.compiler_flags.push(flag.into());
self
}
pub fn host_compiler_flag<F: Into<String>>(&mut self, flag: F) -> &mut Build {
self.compiler_flags
.extend(["-Xcompiler".to_string(), flag.into()]);
self
}
pub fn nvcc_flags<I>(&mut self, flags: I) -> &mut Self
where
I: IntoIterator,
I::Item: Into<String>,
{
for flag in flags {
self.nvcc_flag(flag);
}
self
}
pub fn host_compiler_flags<I>(&mut self, flags: I) -> &mut Self
where
I: IntoIterator,
I::Item: Into<String>,
{
for flag in flags {
self.host_compiler_flag(flag);
}
self
}
pub fn include<P: Into<PathBuf>>(&mut self, dir: P) -> &mut Self {
self.include_directories.push(dir.into());
self
}
pub fn includes<P>(&mut self, dirs: P) -> &mut Self
where
P: IntoIterator,
P::Item: Into<PathBuf>,
{
for dir in dirs {
self.include(dir);
}
self
}
pub fn warnings(&mut self, enable: bool) -> &mut Self {
self.warnings = enable;
self
}
pub fn warnings_as_errors(&mut self, enable: bool) -> &mut Self {
self.warnings_as_errors = enable;
self
}
}