rustytorch/crates/core/rtx-backend-cuda/build.rs

//! Build script for RTX-Backend-CUDA kernels
//!
//! Compiles CUDA kernels to PTX and embeds them in the binary.

use std::env;
use std::fs;
use std::path::Path;
use std::process::Command;

fn main() {
    compile_cuda_kernels();
}

fn compile_cuda_kernels() {
    let out_dir = env::var("OUT_DIR").unwrap();

    // Path to the shared CUDA kernels in rtx-tensor
    let cuda_kernels_dir = Path::new("../rtx-tensor/src/cuda_kernels");
    let element_wise_cu = cuda_kernels_dir.join("element_wise.cu");

    println!("cargo:rerun-if-changed={}", element_wise_cu.display());

    // Check if NVCC is available
    if !is_nvcc_available() {
        println!("cargo:warning=NVCC not found - using dummy PTX (GPU kernels will not work)");
        create_dummy_ptx(&out_dir);
        return;
    }

    println!("cargo:warning=Compiling CUDA kernels for rtx-backend-cuda");

    // Compile element-wise kernels
    if element_wise_cu.exists() {
        compile_kernel(
            &element_wise_cu,
            &Path::new(&out_dir).join("element_wise.ptx"),
        )
        .expect("Failed to compile element-wise CUDA kernels");
        println!("cargo:warning=CUDA kernels compiled successfully");
    } else {
        println!(
            "cargo:warning=CUDA kernel source not found at {:?}, using dummy PTX",
            element_wise_cu
        );
        create_dummy_ptx(&out_dir);
    }
}

fn is_nvcc_available() -> bool {
    Command::new("nvcc")
        .arg("--version")
        .output()
        .map(|output| output.status.success())
        .unwrap_or(false)
}

fn compile_kernel<P: AsRef<Path>>(
    cu_file: P,
    ptx_file: P,
) -> Result<(), Box<dyn std::error::Error>> {
    let cu_path = cu_file.as_ref();
    let ptx_path = ptx_file.as_ref();

    // Ensure output directory exists
    if let Some(parent) = ptx_path.parent() {
        fs::create_dir_all(parent)?;
    }

    // Compile CUDA to PTX
    // Use sm_120 (RTX 5060 Ti / Blackwell, compute 12.0) for the target GPU
    // Use PTX ISA 7.5 for driver compatibility (CUDA 11.x compatible)
    let output = Command::new("nvcc")
        .arg("--ptx") // Generate PTX intermediate representation
        .arg("--gpu-architecture=sm_120") // RTX 5060 Ti / Blackwell (compute 12.0)
        .arg("--allow-unsupported-compiler") // Allow GCC 13+ (needed for Ubuntu 24.04)
        .arg("-Xptxas")
        .arg("--allow-expensive-optimizations=true")
        .arg("--optimize=2") // Maximum optimization
        .arg("--use_fast_math") // Fast math for performance
        .arg("--maxrregcount=64") // Register optimization
        .arg("-o")
        .arg(ptx_path) // Output PTX file
        .arg(cu_path) // Input CUDA source
        .output()?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        eprintln!("NVCC compilation failed:");
        eprintln!("{stderr}");
        return Err(format!("NVCC failed with: {stderr}").into());
    }

    println!(
        "cargo:warning=Compiled {} -> {}",
        cu_path.display(),
        ptx_path.display()
    );
    Ok(())
}

fn create_dummy_ptx(out_dir: &str) {
    // Create dummy PTX file when NVCC is not available
    let dummy_ptx = r"
.version 7.0
.target sm_50
.address_size 64

.visible .entry dummy_kernel(.param .u64 dummy_param) {
    ret;
}
";

    fs::write(Path::new(out_dir).join("element_wise.ptx"), dummy_ptx)
        .expect("Failed to create dummy PTX file");
}