foxhunt/ml/build.rs

use std::env;
use std::path::PathBuf;

#[cfg(feature = "cuda")]
fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("cargo:rerun-if-changed=src/liquid/cuda/liquid_kernels.cu");
    println!("cargo:rerun-if-changed=build.rs");

    println!("cargo:rustc-cfg=cuda_build_enabled");
    if cfg!(feature = "cudnn") {
        println!("cargo:rustc-cfg=cudnn_build_enabled");
    }

    compile_cuda_kernels()?;
    link_cuda_libraries()?;
    setup_cuda_environment()?;

    Ok(())
}

fn compile_cuda_kernels() -> Result<(), Box<dyn std::error::Error>> {
    println!("cargo:info=Compiling CUDA kernels for ML acceleration");
    
    // Check for CUDA compiler
    let nvcc_path = find_nvcc()?;
    println!("cargo:info=Found nvcc at: {}", nvcc_path.display());
    
    // CUDA kernel source path
    let kernel_source = "src/liquid/cuda/liquid_kernels.cu";
    
    // Output directory for compiled objects
    let out_dir = env::var("OUT_DIR")?;
    
    // Compile CUDA kernels
    let mut nvcc_cmd = std::process::Command::new(nvcc_path);
    nvcc_cmd
        .args([
            "--compile",
            "-o", &format!("{}/liquid_kernels.o", out_dir),
            kernel_source,
            "--compiler-options", "-fPIC",
            "--gpu-architecture=sm_75",  // RTX 2060 and newer
            "--gpu-architecture=sm_86",  // RTX 3060 and newer  
            "--gpu-architecture=sm_89",  // RTX 4060 and newer
            "--optimize=3",
            "--use_fast_math",
            "--restrict",
            "--maxrregcount=64",
            "-DCUDA_API_PER_THREAD_DEFAULT_STREAM",
        ]);
    
    println!("cargo:info=Running nvcc command");
    
    let output = nvcc_cmd.output()?;
    
    if !output.status.success() {
        println!("cargo:warning=CUDA kernel compilation failed");
        println!("cargo:warning=STDOUT: {}", String::from_utf8_lossy(&output.stdout));
        println!("cargo:warning=STDERR: {}", String::from_utf8_lossy(&output.stderr));
        return Err("CUDA compilation failed".into());
    }
    
    println!("cargo:info=CUDA kernels compiled successfully");
    
    // Archive the object file
    let lib_path = format!("{}/libliquid_kernels.a", out_dir);
    let mut ar_cmd = std::process::Command::new("ar");
    ar_cmd.args(["rcs", &lib_path, &format!("{}/liquid_kernels.o", out_dir)]);
    
    let ar_output = ar_cmd.output()?;
    if !ar_output.status.success() {
        println!("cargo:warning=Failed to archive CUDA kernels");
        return Err("CUDA archiving failed".into());
    }
    
    // Tell Cargo to link our compiled kernels
    println!("cargo:rustc-link-search=native={}", out_dir);
    println!("cargo:rustc-link-lib=static=liquid_kernels");
    
    Ok(())
}

fn link_cuda_libraries() -> Result<(), Box<dyn std::error::Error>> {
    println!("cargo:info=Setting up CUDA library linking");
    
    // Find CUDA installation
    let cuda_root = find_cuda_root()?;
    println!("cargo:info=CUDA root found at: {}", cuda_root.display());
    
    // CUDA library paths
    let cuda_lib_path = cuda_root.join("lib64");
    let cuda_lib_path_alt = cuda_root.join("lib");
    
    if cuda_lib_path.exists() {
        println!("cargo:rustc-link-search=native={}", cuda_lib_path.display());
    } else if cuda_lib_path_alt.exists() {
        println!("cargo:rustc-link-search=native={}", cuda_lib_path_alt.display());
    } else {
        println!("cargo:warning=No CUDA library path found");
    }
    
    // Essential CUDA libraries for ML acceleration
    println!("cargo:rustc-link-lib=dylib=cuda");           // CUDA driver API
    println!("cargo:rustc-link-lib=dylib=cudart");         // CUDA runtime API
    println!("cargo:rustc-link-lib=dylib=cublas");         // Basic Linear Algebra on CUDA
    println!("cargo:rustc-link-lib=dylib=cublasLt");       // CUDA Basic Linear Algebra LT
    println!("cargo:rustc-link-lib=dylib=curand");         // CUDA Random Number Generation
    println!("cargo:rustc-link-lib=dylib=cufft");          // CUDA Fast Fourier Transform
    
    // Optional: cuDNN for deep learning primitives
    if cfg!(feature = "cudnn") {
        println!("cargo:rustc-link-lib=dylib=cudnn");      // NVIDIA cuDNN
    }
    
    // Optional: NCCL for multi-GPU communication
    if let Ok(nccl_path) = env::var("NCCL_ROOT") {
        let nccl_lib = PathBuf::from(nccl_path).join("lib");
        if nccl_lib.exists() {
            println!("cargo:rustc-link-search=native={}", nccl_lib.display());
            println!("cargo:rustc-link-lib=dylib=nccl");
        }
    }
    
    Ok(())
}

fn setup_cuda_environment() -> Result<(), Box<dyn std::error::Error>> {
    println!("cargo:info=Setting up CUDA environment variables");
    
    // Set CUDA-specific flags
    println!("cargo:rustc-cfg=cuda_enabled");
    
    // GPU architecture defines
    println!("cargo:rustc-cfg=gpu_sm_75"); // RTX 2060+
    println!("cargo:rustc-cfg=gpu_sm_86"); // RTX 3060+
    println!("cargo:rustc-cfg=gpu_sm_89"); // RTX 4060+
    
    // CUDA API version detection
    let cuda_version = detect_cuda_version()?;
    println!("cargo:rustc-cfg=cuda_version_major=\"{}\"", cuda_version.0);
    println!("cargo:rustc-cfg=cuda_version_minor=\"{}\"", cuda_version.1);
    
    if cuda_version.0 >= 12 {
        println!("cargo:rustc-cfg=cuda_12_plus");
    }
    if cuda_version.0 >= 11 {
        println!("cargo:rustc-cfg=cuda_11_plus");
    }
    
    Ok(())
}

fn find_nvcc() -> Result<PathBuf, Box<dyn std::error::Error>> {
    // Try multiple common locations for nvcc
    let candidates = [
        "/usr/local/cuda/bin/nvcc",
        "/usr/local/cuda-12.0/bin/nvcc", 
        "/usr/local/cuda-11.8/bin/nvcc",
        "/usr/local/cuda-11.7/bin/nvcc",
        "/opt/cuda/bin/nvcc",
        "nvcc", // In PATH
    ];
    
    // Check environment variable first
    if let Ok(cuda_home) = env::var("CUDA_HOME") {
        let nvcc_path = PathBuf::from(cuda_home).join("bin/nvcc");
        if nvcc_path.exists() {
            return Ok(nvcc_path);
        }
    }
    
    if let Ok(cuda_root) = env::var("CUDA_ROOT") {
        let nvcc_path = PathBuf::from(cuda_root).join("bin/nvcc");
        if nvcc_path.exists() {
            return Ok(nvcc_path);
        }
    }
    
    // Try common paths
    for candidate in &candidates {
        let path = PathBuf::from(candidate);
        if path.exists() {
            return Ok(path);
        }
    }
    
    // Try which/where command
    if let Ok(output) = std::process::Command::new("which").arg("nvcc").output() {
        if output.status.success() {
            let path_str = String::from_utf8_lossy(&output.stdout);
            let path_str_trimmed = path_str.trim().to_owned();
            if !path_str_trimmed.is_empty() {
                return Ok(PathBuf::from(path_str_trimmed));
            }
        }
    }
    
    Err("nvcc (NVIDIA CUDA Compiler) not found. Please install CUDA toolkit or set CUDA_HOME environment variable.".into())
}

fn find_cuda_root() -> Result<PathBuf, Box<dyn std::error::Error>> {
    // Try environment variables
    if let Ok(cuda_home) = env::var("CUDA_HOME") {
        let path = PathBuf::from(cuda_home);
        if path.exists() {
            return Ok(path);
        }
    }
    
    if let Ok(cuda_root) = env::var("CUDA_ROOT") {
        let path = PathBuf::from(cuda_root);
        if path.exists() {
            return Ok(path);
        }
    }
    
    // Try common installation paths
    let candidates = [
        "/usr/local/cuda",
        "/usr/local/cuda-12.0",
        "/usr/local/cuda-11.8", 
        "/usr/local/cuda-11.7",
        "/opt/cuda",
    ];
    
    for candidate in &candidates {
        let path = PathBuf::from(candidate);
        if path.exists() {
            return Ok(path);
        }
    }
    
    Err("CUDA installation not found. Please install CUDA toolkit or set CUDA_HOME.".into())
}

fn detect_cuda_version() -> Result<(u32, u32), Box<dyn std::error::Error>> {
    let nvcc_path = find_nvcc()?;
    
    let output = std::process::Command::new(nvcc_path)
        .arg("--version")
        .output()?;
    
    if !output.status.success() {
        return Err("Failed to get CUDA version".into());
    }
    
    let version_text = String::from_utf8_lossy(&output.stdout);
    
    // Parse version from nvcc output
    // Example: "Cuda compilation tools, release 12.0, V12.0.140"
    for line in version_text.lines() {
        if line.contains("release") {
            if let Some(version_part) = line.split("release ").nth(1) {
                if let Some(version_str) = version_part.split(',').next() {
                    let parts: Vec<&str> = version_str.split('.').collect();
                    if parts.len() >= 2 {
                        let major: u32 = parts.first().and_then(|s| s.parse().ok()).unwrap_or(11);
                        let minor: u32 = parts.get(1).and_then(|s| s.parse().ok()).unwrap_or(0);
                        return Ok((major, minor));
                    }
                }
            }
        }
    }
    
    println!("cargo:warning=Could not parse CUDA version, assuming 11.0");
    Ok((11, 0))
}

#[cfg(not(feature = "cuda"))]
fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("cargo:rerun-if-changed=build.rs");
    println!("cargo:info=CUDA support disabled - building CPU-only version");
    println!("cargo:rustc-cfg=cpu_only_build");
    Ok(())
}