lufloat 0.1.2

Fastest FP16 Math and AI Library for AMD APUs
Documentation
use std::{
    env::var,
    fs::read_dir,
    path::PathBuf,
    process::{Child, Command},
};

fn main() {
    if var("DOCS_RS").is_ok() {
        return;
    }
    println!("cargo:rerun-if-changed=kernels");
    println!("cargo:rerun-if-env-changed=ROCM_PATH");
    println!("cargo:rerun-if-env-changed=HIPCC");
    println!("cargo:rerun-if-env-changed=HIP_PATH");
    let out_dir = var("OUT_DIR").unwrap();
    let collected = compile_lib(&out_dir);
    let mut obj_files = Vec::with_capacity(collected.len());
    for (mut child, obj_file) in collected {
        assert!(child.wait().unwrap().success());
        obj_files.push(obj_file);
    }
    static_archive_lib(&out_dir, &obj_files);
    link_lib(&out_dir);
}

fn compile_lib(out_dir: &str) -> Vec<(Child, PathBuf)> {
    read_dir("kernels")
        .unwrap()
        .map(|entry| {
            let path = entry.unwrap().path();
            let display = path.display();
            println!("cargo:rerun-if-changed={display}");
            let name = path.file_name().unwrap();
            let obj = PathBuf::from(out_dir).join(name).with_extension("o");
            let child = Command::new("hipcc")
                .arg("-c")
                .arg(&path)
                .arg("-o")
                .arg(&obj)
                .args([
                    "-Ofast",
                    "-fPIC",
                    "-fgpu-flush-denormals-to-zero",
                    "--gpu-max-threads-per-block=256",
                    "-munsafe-fp-atomics",
                    "--offload-arch=native",
                ])
                .spawn()
                .unwrap();
            (child, obj)
        })
        .collect()
}

fn static_archive_lib(out_dir: &str, obj_files: &[PathBuf]) {
    assert!(
        Command::new("ar")
            .arg("rcs")
            .arg(PathBuf::from(out_dir).join("libkernels.a"))
            .args(obj_files)
            .status()
            .unwrap()
            .success()
    );
}

fn link_lib(out_dir: &str) {
    let rocm_hip_path = var("HIP_PATH")
        .or_else(|_| var("ROCM_PATH"))
        .unwrap_or_else(|_| String::from("/opt/rocm"));
    println!("cargo:rustc-link-search=native={out_dir}");
    println!("cargo:rustc-link-lib=static=kernels");
    println!("cargo:rustc-link-search=native={rocm_hip_path}/lib");
    println!("cargo:rustc-link-lib=dylib=amdhip64");
}