cubecl_hip_sys/
dynamic.rs

1use libloading::{Library, Symbol};
2use std::{env, path::PathBuf, sync::OnceLock};
3
4struct Libraries {
5    hip: Library,
6    hiprtc: Library,
7}
8
9static LIBRARIES: OnceLock<Result<Libraries, String>> = OnceLock::new();
10
11/// Returns whether both HIP runtime libraries can be loaded.
12pub fn is_available() -> bool {
13    libraries().is_ok()
14}
15
16/// Resolve a HIP symbol without making the HIP libraries link-time dependencies.
17///
18/// This function is called by the generated bindings. A missing runtime is a
19/// runtime error because the public binding functions cannot return one common
20/// error type for all of their C signatures.
21pub(crate) unsafe fn load<T: Copy>(name: &[u8]) -> T {
22    let libraries = libraries()
23        .as_ref()
24        .unwrap_or_else(|error| panic!("{error}"));
25    let library = if name.starts_with(b"hiprtc") {
26        &libraries.hiprtc
27    } else {
28        &libraries.hip
29    };
30
31    let symbol: Symbol<'_, T> = unsafe { library.get(name) }.unwrap_or_else(|error| {
32        let symbol_name = name.strip_suffix(&[0]).unwrap_or(name);
33        let symbol = String::from_utf8_lossy(symbol_name);
34        panic!("HIP symbol `{symbol}` is unavailable: {error}");
35    });
36    *symbol
37}
38
39fn libraries() -> &'static Result<Libraries, String> {
40    LIBRARIES.get_or_init(|| unsafe { load_libraries() })
41}
42
43unsafe fn load_libraries() -> Result<Libraries, String> {
44    let search_paths = search_paths();
45    let hip = load_library("amdhip64", &search_paths)?;
46    let hiprtc = load_library("hiprtc", &search_paths)?;
47    Ok(Libraries { hip, hiprtc })
48}
49
50fn search_paths() -> Vec<PathBuf> {
51    ["ROCM_PATH", "HIP_PATH"]
52        .into_iter()
53        .filter_map(env::var_os)
54        .flat_map(|path| {
55            let path = PathBuf::from(path);
56            [path.join("lib"), path]
57        })
58        .collect()
59}
60
61unsafe fn load_library(name: &str, search_paths: &[PathBuf]) -> Result<Library, String> {
62    let names = library_names(name);
63    let mut errors = Vec::new();
64
65    for path in search_paths {
66        for library_name in &names {
67            let candidate = path.join(library_name);
68            match unsafe { Library::new(&candidate) } {
69                Ok(library) => return Ok(library),
70                Err(error) => errors.push(format!("{}: {error}", candidate.display())),
71            }
72        }
73    }
74
75    for library_name in &names {
76        match unsafe { Library::new(library_name) } {
77            Ok(library) => return Ok(library),
78            Err(error) => errors.push(format!("{library_name}: {error}")),
79        }
80    }
81
82    Err(format!(
83        "Could not load HIP library `{name}`. Install ROCm or set ROCM_PATH/HIP_PATH.\n{}",
84        errors.join("\n")
85    ))
86}
87
88fn library_names(name: &str) -> Vec<String> {
89    if cfg!(target_os = "windows") {
90        vec![format!("{name}.dll")]
91    } else if cfg!(target_os = "macos") {
92        vec![format!("lib{name}.dylib")]
93    } else {
94        vec![
95            format!("lib{name}.so"),
96            format!("lib{name}.so.1"),
97            format!("lib{name}.so.0"),
98        ]
99    }
100}
cubecl_hip_sys/dynamic.rs

cubecl_hip_sys/
dynamic.rs
