torch-build

用于链接 libtorch FFI 接口的工具。

使用方法

在您的 Cargo.toml 中添加 cc 和 torch-build 构建依赖项。

[build-dependencies]
anyhow = "1.0.45"
cc = "1.0.72"
torch-build = "X.Y.Z"

build.rs 代码片段编译 C++/CUDA 源文件并将其链接到 libtorch。

use anyhow::Result;

// Provides C++/CUDA source files
const CPP_SOURCE: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/csrc/nms_cpu.cpp");
const CUDA_SOURCE: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/csrc/nms_cuda.cu");

fn main() -> Result<()> {
    // The container stores cargo-build commands
    let mut cargo_commands = vec![];

    // Compile C++ files
    let mut build = cc::Build::new();
    torch_build::build_cpp(
        &mut build,
        true,
        false,
        Some(&mut cargo_commands),
        [CPP_SOURCE],
    )?;
    build.try_compile("nms_cpu")?;

    // Compile CUDA files
    let mut build = cc::Build::new();
    torch_build::build_cuda(&mut build, false, Some(&mut cargo_commands), [CUDA_SOURCE])?;
    build.try_compile("nms_cuda")?;

    // Re-compile if C++/CUDA source files were modified
    println!("cargo:rerun-if-changed={}", CPP_SOURCE);
    println!("cargo:rerun-if-changed={}", CUDA_SOURCE);
    cargo_commands.iter().for_each(|command| {
        println!("{}", command);
    });

    Ok(())
}

许可证

MIT 许可证。请参阅 许可证文件。