rustytorch/crates/training/rtx-compress/build.rs

//! Build script for RTX-Compress CUDA kernels
//!
//! Compiles MX quantization CUDA kernels to PTX for GPU acceleration.

use std::env;

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

    // Avoid unused warnings when cuda feature is not enabled
    #[cfg(not(feature = "cuda"))]
    {
        let _ = env::var("OUT_DIR");
    }
}

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

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

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

    println!("cargo:warning=Compiling MX quantization CUDA kernels");

    // Compile MX kernels
    match compile_kernel(
        cuda_kernels_dir.join("mx_kernels.cu"),
        Path::new(&out_dir).join("mx_kernels.ptx"),
    ) {
        Ok(_) => {
            println!("cargo:warning=MX CUDA kernels compiled successfully");
        }
        Err(e) => {
            println!("cargo:warning=Failed to compile MX CUDA kernels: {e}");
            println!("cargo:warning=MX GPU quantization will fall back to CPU");
            create_dummy_ptx(&out_dir);
        }
    }
}

#[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
    // This allows the crate to compile but GPU kernels won't work
    let dummy_ptx = r#"
.version 7.0
.target sm_86
.address_size 64

// Dummy kernel - real implementation requires NVCC compilation
.visible .entry mx_compute_exponents_kernel(
    .param .u64 input,
    .param .u64 exponents,
    .param .u64 numel,
    .param .u32 format
) {
    ret;
}

.visible .entry mx_quantize_mantissas_8bit_kernel(
    .param .u64 input,
    .param .u64 exponents,
    .param .u64 mantissas,
    .param .u64 numel,
    .param .u32 format
) {
    ret;
}

.visible .entry mx_quantize_mantissas_4bit_kernel(
    .param .u64 input,
    .param .u64 exponents,
    .param .u64 mantissas,
    .param .u64 numel,
    .param .u32 format
) {
    ret;
}

.visible .entry mx_dequantize_8bit_kernel(
    .param .u64 mantissas,
    .param .u64 exponents,
    .param .u64 output,
    .param .u64 numel,
    .param .u32 format
) {
    ret;
}

.visible .entry mx_dequantize_4bit_kernel(
    .param .u64 mantissas,
    .param .u64 exponents,
    .param .u64 output,
    .param .u64 numel,
    .param .u32 format
) {
    ret;
}

.visible .entry mx_fused_quantize_8bit_kernel(
    .param .u64 input,
    .param .u64 exponents,
    .param .u64 mantissas,
    .param .u64 numel,
    .param .u32 format
) {
    ret;
}

.visible .entry mx_fused_quantize_4bit_kernel(
    .param .u64 input,
    .param .u64 exponents,
    .param .u64 mantissas,
    .param .u64 numel,
    .param .u32 format
) {
    ret;
}
"#;

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