hipster/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
use std::{
    env::{self, var, VarError},
    path::PathBuf,
};

const DEFAULT_OPENCL_PATH: &str = "/usr";

// fn main() {
//     let paths = find_opencl_root_dirs();
//     if paths.is_empty() {
//         panic!("Could not find an OpenCL installation.")
//     }

//     for path in paths.iter() {
//         println!(
//             "cargo:rustc-link-search=native={}",
//             path.join("lib").display()
//         );
//     }
//     println!("cargo:rustc-link-lib=dylib=OpenCL");
//     println!("cargo:rerun-if-changed=build.rs");
//     println!("cargo:rerun-if-env-changed=OPENCL_LIBRARY_PATH");
// }

pub fn include_hip() {
    if env::var("DOCS_RS").is_err() && !cfg!(doc) {
        let paths = find_hip_root_dirs();
        if paths.is_empty() {
            panic!("Could not find a hip installation");
        }
        for path in paths {
            println!("cargo:rustc-link-search=native={}/lib", path.display());
        }

        println!("cargo:rustc-link-lib=dylib=amdhip64");
        println!("cargo:rerun-if-changed=build.rs");
        println!("cargo:rerun-if-env-changed=CUDA_LIBRARY_PATH");
        println!("cargo:rerun-if-env-changed=CUDA_ROOT");
        println!("cargo:rerun-if-env-changed=CUDA_PATH");
        println!("cargo:rerun-if-env-changed=CUDA_TOOLKIT_ROOT_DIR");
        println!("cargo:rerun-if-env-changed=HIP_PATH");
        println!("cargo:rerun-if-env-changed=HCC_ROCCLR_HOME");
        println!("cargo:rerun-if-env-changed=HSA_PATH");
        println!("cargo:rerun-if-env-changed=ROCM_PATH");
        // println!("cargo:rerun-if-env-changed=HIP_CLANG_PATH");
        // println!("cargo:rerun-if-env-changed=DEVICE_LIB_PATH");
    }
}

pub fn read_env() -> Vec<PathBuf> {
    let split_char = if cfg!(target_os = "windows") {
        ";"
    } else {
        ":"
    };
    // Link to OpenCL Headers. If HIP_PATH isn't supplied, use a default.
    match var("OPENCL_LIBRARY_PATH") {
        // The location of the libopencl can be hardcoded with the
        // OPENCL_LIBRARY_PATH environment variable.
        Ok(path) => {
            path.split(split_char)
                .map(PathBuf::from)
                // Check that the path exists.
                .filter(|pb| pb.exists())
                .collect()
        }
        Err(VarError::NotPresent) => DEFAULT_OPENCL_PATH
            .split(split_char)
            .map(PathBuf::from)
            // Check that the path exists.
            .filter(|pb| pb.exists())
            .collect(),
        Err(e @ VarError::NotUnicode(_)) => panic!("{}: HIP_PATH: {}", env!("CARGO_PKG_NAME"), e),
    }
}

pub fn find_hip_root_dirs() -> Vec<PathBuf> {
    let mut candidates = read_env();
    candidates.push(PathBuf::from("/opt/cuda"));
    candidates.push(PathBuf::from("/opt/cuda/targets/x86_64-linux"));
    candidates.push(PathBuf::from("/opt/rocm"));
    candidates.push(PathBuf::from("/opt/rocm/opencl"));
    candidates.push(PathBuf::from("/usr"));

    let mut valid_paths = vec![];
    for base in &candidates {
        let header = base.join("include/CL/cl.h");
        let lib = PathBuf::from(base).join("lib");
        if lib.is_dir() && header.is_file() {
            valid_paths.push(base.clone());
        }
    }
    valid_paths
}