rustytorch/crates/training/rtx-flash-attention/build.rs

//! Build script for Flash Attention CUDA kernels.
//!
//! Compute capability is selected at build time:
//!   - Set `CUDA_ARCH` env var (e.g. `CUDA_ARCH=sm_120`) to override.
//!   - Default: `sm_120` (Blackwell — RTX 5000 series).
//!
//! `CUDA_PATH` must point to the CUDA toolkit root if `nvcc` is not in `$PATH`
//! (e.g. `CUDA_PATH=/usr/local/cuda-13.1`).

use std::env;
use std::path::{Path, PathBuf};
use std::process::Command;

fn main() {
    println!("cargo:rerun-if-changed=cuda/");
    println!("cargo:rerun-if-changed=metal/");
    println!("cargo:rerun-if-changed=build.rs");

    let cuda_feature = env::var("CARGO_FEATURE_CUDA").is_ok();
    let metal_feature = env::var("CARGO_FEATURE_METAL").is_ok();

    if cuda_feature {
        if let Some(nvcc) = find_nvcc() {
            println!("cargo:rustc-cfg=cuda_available");
            compile_cuda_kernels(&nvcc);
            link_cuda_libraries();
        } else {
            println!(
                "cargo:warning=CUDA feature enabled but nvcc not found. Set CUDA_PATH=/path/to/cuda"
            );
        }
    }

    if metal_feature {
        println!("cargo:rustc-cfg=metal_available");
        compile_metal_shaders();
    }
}

/// Locate the `nvcc` binary: check `$PATH` first, then `$CUDA_PATH/bin/nvcc`.
fn find_nvcc() -> Option<PathBuf> {
    // 1. Try bare name (nvcc in PATH).
    if let Ok(output) = Command::new("nvcc").arg("--version").output() {
        if output.status.success() {
            return Some(PathBuf::from("nvcc"));
        }
    }

    // 2. Try CUDA_PATH / CUDA_HOME env vars.
    for var in &["CUDA_PATH", "CUDA_HOME"] {
        if let Ok(root) = env::var(var) {
            let nvcc = Path::new(&root).join("bin").join("nvcc");
            if nvcc.exists() {
                return Some(nvcc);
            }
        }
    }

    // 3. Try common installation paths.
    for prefix in &[
        "/usr/local/cuda",
        "/usr/local/cuda-13.1",
        "/usr/local/cuda-12",
        "/opt/cuda",
    ] {
        let nvcc = Path::new(prefix).join("bin").join("nvcc");
        if nvcc.exists() {
            return Some(nvcc);
        }
    }

    None
}

/// Pick the target SM from the `CUDA_ARCH` env var, defaulting to `sm_120`
/// (Blackwell — RTX 5060 Ti / RTX 5090).
fn cuda_arch() -> (String, String) {
    let arch = env::var("CUDA_ARCH").unwrap_or_else(|_| "sm_120".to_owned());
    // arch is e.g. "sm_120"; compute is "compute_120".
    let compute = arch.replace("sm_", "compute_");
    (compute, arch)
}

fn compile_cuda_kernels(nvcc: &Path) {
    let out_dir = env::var("OUT_DIR").expect("OUT_DIR not set");
    let cuda_dir = PathBuf::from("cuda");
    let (compute, sm) = cuda_arch();

    let cuda_flags = [
        format!("--gpu-architecture={compute}"),
        format!("--gpu-code={sm}"),
        "--use_fast_math".to_owned(),
        "--maxrregcount=255".to_owned(),
        "-O3".to_owned(),
        "--extra-device-vectorization".to_owned(),
        "--fmad=true".to_owned(),
        "--prec-div=false".to_owned(),
        "--prec-sqrt=false".to_owned(),
        "--ftz=true".to_owned(),
    ];

    let kernels = [
        "flash_attention_forward.cu",
        "flash_attention_backward.cu",
        "online_softmax.cu",
        "utils.cu",
    ];

    for kernel in &kernels {
        let input_path = cuda_dir.join(kernel);
        let output_name = kernel.replace(".cu", ".ptx");
        let output_path = Path::new(&out_dir).join(&output_name);

        let mut cmd = Command::new(nvcc);
        cmd.arg("--ptx")
            .arg(&input_path)
            .arg("-o")
            .arg(&output_path);
        for flag in &cuda_flags {
            cmd.arg(flag);
        }

        // Include CUDA headers.
        for var in &["CUDA_PATH", "CUDA_HOME"] {
            if let Ok(root) = env::var(var) {
                cmd.arg(format!("-I{root}/include"));
                break;
            }
        }

        let output = cmd
            .output()
            .unwrap_or_else(|e| panic!("nvcc exec failed for {kernel}: {e}"));
        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            panic!("CUDA kernel compilation failed for {kernel}:\n{stderr}");
        }
    }

    combine_ptx_files(&out_dir, &sm);
}

fn combine_ptx_files(out_dir: &str, sm: &str) {
    use std::fs::{File, OpenOptions};
    use std::io::{BufWriter, Read, Write};

    let combined_path = Path::new(out_dir).join("flash_attention_kernels.ptx");
    let mut out = BufWriter::new(
        OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(&combined_path)
            .expect("Failed to create combined PTX file"),
    );

    // Derive PTX version from SM: sm_120 → .version 8.0 target sm_120.
    // PTX ISA 8.0 covers sm_100+ (Blackwell).
    let ptx_version = if sm >= "sm_100" { "8.0" } else { "7.8" };
    writeln!(out, ".version {ptx_version}").unwrap();
    writeln!(out, ".target {sm}").unwrap();
    writeln!(out, ".address_size 64").unwrap();
    writeln!(out).unwrap();

    for ptx_file in &[
        "flash_attention_forward.ptx",
        "flash_attention_backward.ptx",
        "online_softmax.ptx",
        "utils.ptx",
    ] {
        let ptx_path = Path::new(out_dir).join(ptx_file);
        if ptx_path.exists() {
            let mut contents = String::new();
            File::open(&ptx_path)
                .expect("Failed to open PTX file")
                .read_to_string(&mut contents)
                .expect("Failed to read PTX file");

            // Skip the per-file header and append body.
            let skip = contents
                .lines()
                .position(|l| l.is_empty())
                .unwrap_or(0)
                .saturating_add(1);
            for line in contents.lines().skip(skip) {
                writeln!(out, "{line}").unwrap();
            }
            writeln!(out).unwrap();
        }
    }

    out.flush().unwrap();
    drop(out);

    println!(
        "cargo:rustc-env=FLASH_ATTENTION_PTX_PATH={}",
        combined_path.display()
    );
}

fn link_cuda_libraries() {
    // Search in CUDA_PATH, CUDA_HOME, or well-known locations.
    let lib_found = ["CUDA_PATH", "CUDA_HOME"]
        .iter()
        .find_map(|var| env::var(var).ok())
        .map(|root| {
            println!("cargo:rustc-link-search=native={root}/lib64");
            true
        })
        .unwrap_or(false);

    if !lib_found {
        for path in &[
            "/usr/local/cuda/lib64",
            "/usr/local/cuda-13.1/lib64",
            "/opt/cuda/lib64",
            "/usr/lib/x86_64-linux-gnu",
        ] {
            if Path::new(path).exists() {
                println!("cargo:rustc-link-search=native={path}");
                break;
            }
        }
    }

    for lib in &["cudart", "cublas", "curand", "cusparse", "cufft"] {
        println!("cargo:rustc-link-lib={lib}");
    }
}

fn compile_metal_shaders() {
    let metal_dir = PathBuf::from("metal");
    let shader = metal_dir.join("flash_attention.metal");
    if shader.exists() {
        let abs = std::fs::canonicalize(&shader).expect("canonicalize metal shader");
        println!("cargo:rustc-env=METAL_SHADER_PATH={}", abs.display());
    }
}