rustytorch/crates/core/rtx-tensor/build.rs

//! Build script for RTX-Tensor CUDA kernels
//!
//! Compiles CUDA kernels to PTX and embeds them in the binary
//! for runtime kernel loading and execution.

#[allow(unused_imports)]
use std::env;
#[allow(unused_imports)]
use std::fs;
#[allow(unused_imports)]
use std::path::Path;
#[allow(unused_imports)]
use std::process::Command;

fn main() {
    // Only link OpenBLAS if the feature is enabled
    #[cfg(feature = "openblas")]
    {
        // Force proper OpenBLAS linking configuration
        // The openblas-src crate builds the library but linking may fail
        // We need to explicitly tell the linker where to find it

        // Get the OpenBLAS library path from the build output
        if let Ok(out_dir) = env::var("DEP_OPENBLAS_LIB") {
            println!("cargo:rustc-link-search=native={out_dir}");
            println!("cargo:rustc-link-lib=static=openblas");
        } else {
            // Fallback: let openblas-src handle linking automatically
            println!("cargo:rustc-link-lib=openblas");
        }

        // Ensure CBLAS symbols are available
        println!("cargo:rustc-link-lib=gfortran"); // OpenBLAS may need Fortran runtime
    }

    // Only compile CUDA kernels if CUDA feature is enabled
    #[cfg(feature = "cuda")]
    compile_cuda_kernels();

    // Link cuSPARSELt if the feature is enabled
    #[cfg(feature = "cusparselt_link")]
    link_cusparselt();
}

/// Configure cuSPARSELt library linking
#[cfg(feature = "cusparselt_link")]
fn link_cusparselt() {
    // Check common paths for libcusparseLt
    let lib_paths = [
        "/usr/lib/x86_64-linux-gnu/libcusparseLt/13",
        "/usr/lib/x86_64-linux-gnu",
        "/usr/local/cuda/lib64",
        "/usr/local/cuda-13.0/lib64",
        "/usr/local/cuda-13.1/lib64",
    ];

    let mut found = false;
    for path in lib_paths {
        let lib_path = Path::new(path);
        if lib_path.join("libcusparseLt.so").exists() {
            println!("cargo:rustc-link-search=native={}", path);
            found = true;
            break;
        }
    }

    // Also check CUSPARSELT_PATH environment variable
    if !found {
        if let Ok(path) = env::var("CUSPARSELT_PATH") {
            println!("cargo:rustc-link-search=native={}", path);
            found = true;
        }
    }

    // Also check LD_LIBRARY_PATH
    if !found {
        if let Ok(ld_path) = env::var("LD_LIBRARY_PATH") {
            for dir in ld_path.split(':') {
                if Path::new(dir).join("libcusparseLt.so").exists() {
                    println!("cargo:rustc-link-search=native={}", dir);
                    found = true;
                    break;
                }
            }
        }
    }

    if found {
        println!("cargo:rustc-link-lib=dylib=cusparseLt");
        println!("cargo:warning=cuSPARSELt library found and linked for 2:4 sparsity acceleration");
    } else {
        println!("cargo:warning=cuSPARSELt library not found - 2:4 sparsity will fall back to CPU");
    }
}

#[cfg(feature = "cuda")]
fn compile_cuda_kernels() {
    let out_dir = env::var("OUT_DIR").unwrap();
    let cuda_kernels_dir = Path::new("src/cuda_kernels");

    println!("cargo:rerun-if-changed=src/cuda_kernels/element_wise.cu");
    println!("cargo:rerun-if-changed=src/cuda_kernels/fp8_cast.cu");

    // Check if NVCC is available
    if !is_nvcc_available() {
        println!("cargo:warning=NVCC not found - CUDA kernels will not be compiled");
        println!("cargo:warning=GPU acceleration will not be available");
        create_dummy_ptx(&out_dir);
        return;
    }

    println!("cargo:warning=Compiling CUDA kernels for RTX-Tensor GPU acceleration");

    // Compile element-wise kernels
    compile_kernel(
        cuda_kernels_dir.join("element_wise.cu"),
        Path::new(&out_dir).join("element_wise.ptx"),
    )
    .expect("Failed to compile element-wise CUDA kernels");

    // Compile FP8 cast kernels (SM_89+ hardware path; software fallback for older GPUs)
    compile_kernel(
        cuda_kernels_dir.join("fp8_cast.cu"),
        Path::new(&out_dir).join("fp8_cast.ptx"),
    )
    .expect("Failed to compile FP8 cast CUDA kernels");

    println!("cargo:warning=CUDA kernels compiled successfully");
}

#[cfg(feature = "cuda")]
fn is_nvcc_available() -> bool {
    Command::new("nvcc")
        .arg("--version")
        .output()
        .map(|output| output.status.success())
        .unwrap_or(false)
}

#[cfg(feature = "cuda")]
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 with RTX 3050/3090 (Ampere) optimizations
    let output = Command::new("nvcc")
        .arg("--ptx") // Generate PTX intermediate representation
        .arg("--gpu-architecture=sm_86") // RTX 3050/3090 (Ampere) compute capability
        .arg("--allow-unsupported-compiler") // Allow GCC 13+ (needed for Ubuntu 24.04)
        .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(())
}

#[cfg(feature = "cuda")]
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 element_wise PTX file");

    // Also create a dummy fp8_cast.ptx so that `include_str!` in fp8_cast.rs
    // compiles cleanly even when NVCC is unavailable.
    fs::write(Path::new(out_dir).join("fp8_cast.ptx"), dummy_ptx)
        .expect("Failed to create dummy fp8_cast PTX file");
}

/// Helper to embed PTX as string literal at compile time
#[cfg(feature = "cuda")]
pub fn generate_ptx_include() {
    let out_dir = env::var("OUT_DIR").unwrap();
    let ptx_path = Path::new(&out_dir).join("element_wise.ptx");

    if ptx_path.exists() {
        let ptx_content = fs::read_to_string(&ptx_path).expect("Failed to read compiled PTX file");

        let include_code = format!(
            "/// Compiled CUDA kernels as PTX strings\npub mod ptx_kernels {{\n    pub const ELEMENT_WISE_PTX: &str = r#\"{ptx_content}\"#;\n}}"
        );

        let include_path = Path::new(&out_dir).join("ptx_include.rs");
        fs::write(include_path, include_code).expect("Failed to generate PTX include file");
    }
}