micromeasure 0.10.0

Microbenchmark harness for tiny operations and PMU-aware measurement
Documentation
use std::path::{Path, PathBuf};

fn main() {
    println!("cargo:rerun-if-env-changed=CUDA_HOME");
    println!("cargo:rerun-if-env-changed=CUDA_PATH");

    let cuda_enabled = std::env::var_os("CARGO_FEATURE_CUDA").is_some();
    let gpu_counters_enabled = std::env::var_os("CARGO_FEATURE_GPU_COUNTERS").is_some();

    if !cuda_enabled {
        return;
    }

    for lib_dir in cuda_lib_dirs() {
        if lib_dir.exists() {
            println!("cargo:rustc-link-search=native={}", lib_dir.display());
        }
    }

    if gpu_counters_enabled {
        build_gpu_counters();
        println!("cargo:rustc-link-lib=static=micromeasure_gpu_counters");
        println!("cargo:rustc-link-lib=dylib=stdc++");
        println!("cargo:rustc-link-lib=dylib=cuda");
        println!("cargo:rustc-link-lib=dylib=cupti");
        println!("cargo:rerun-if-changed=native/gpu_counters.cpp");
    }
}

fn cuda_lib_dirs() -> Vec<PathBuf> {
    let mut roots = Vec::new();
    if let Some(cuda_home) = std::env::var_os("CUDA_HOME") {
        roots.push(PathBuf::from(cuda_home));
    }
    if let Some(cuda_path) = std::env::var_os("CUDA_PATH") {
        roots.push(PathBuf::from(cuda_path));
    }
    roots.push(PathBuf::from("/usr/local/cuda"));
    roots.push(PathBuf::from("/usr/local/cuda-13"));
    roots.push(PathBuf::from("/usr/local/cuda-13.0"));

    let mut dirs = Vec::new();
    for root in roots {
        push_cuda_lib_dirs(&mut dirs, &root);
    }
    dirs
}

fn push_cuda_lib_dirs(dirs: &mut Vec<PathBuf>, root: &Path) {
    for relative in [
        "lib64",
        "lib",
        "extras/CUPTI/lib64",
        "extras/CUPTI/lib",
        "targets/x86_64-linux/lib",
        "targets/sbsa-linux/lib",
        "targets/aarch64-linux/lib",
    ] {
        let dir = root.join(relative);
        if !dirs.contains(&dir) {
            dirs.push(dir);
        }
    }
}

fn build_gpu_counters() {
    let out_dir = PathBuf::from(std::env::var_os("OUT_DIR").expect("OUT_DIR is set by Cargo"));
    let object = out_dir.join("gpu_counters.o");
    let archive = out_dir.join("libmicromeasure_gpu_counters.a");

    let include_dir = cuda_include_dirs()
        .into_iter()
        .find(|dir| dir.join("cupti_profiler_host.h").exists())
        .expect("CUPTI profiler headers not found; set CUDA_HOME or CUDA_PATH");

    let status = std::process::Command::new("g++")
        .args(["-std=c++17", "-O2", "-I"])
        .arg(&include_dir)
        .args(["-c", "native/gpu_counters.cpp", "-o"])
        .arg(&object)
        .status()
        .expect("failed to run g++ for GPU counters");
    assert!(status.success(), "g++ failed to build GPU counters");

    let status = std::process::Command::new("ar")
        .arg("crs")
        .arg(&archive)
        .arg(&object)
        .status()
        .expect("failed to run ar for GPU counters");
    assert!(status.success(), "ar failed to archive GPU counters");

    println!("cargo:rustc-link-search=native={}", out_dir.display());
}

fn cuda_include_dirs() -> Vec<PathBuf> {
    let mut roots = Vec::new();
    if let Some(cuda_home) = std::env::var_os("CUDA_HOME") {
        roots.push(PathBuf::from(cuda_home));
    }
    if let Some(cuda_path) = std::env::var_os("CUDA_PATH") {
        roots.push(PathBuf::from(cuda_path));
    }
    roots.push(PathBuf::from("/usr/local/cuda"));
    roots.push(PathBuf::from("/usr/local/cuda-13"));
    roots.push(PathBuf::from("/usr/local/cuda-13.0"));

    let mut dirs = Vec::new();
    for root in roots {
        for relative in [
            "include",
            "targets/x86_64-linux/include",
            "targets/sbsa-linux/include",
            "targets/aarch64-linux/include",
        ] {
            let dir = root.join(relative);
            if !dirs.contains(&dir) {
                dirs.push(dir);
            }
        }
    }
    dirs
}