rocmrc 0.1.0

Minimal safe ROCm bindings (HIP, hipRTC), modeled after cudarc
Documentation 
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

//! Minimal safe HIP/ROCm bindings, modeled after `cudarc`.
//!
//! Surface intentionally small: only what `luminal_rocm_lite` needs to host
//! a kernel-launch-based runtime. Grow this crate by following `cudarc`'s
//! API shape so call sites in luminal stay near-identical.

pub mod driver;
pub mod hiprtc;

#[cfg(feature = "rocblas")]
pub mod rocblas;
#[cfg(feature = "hipblaslt")]
pub mod hipblaslt;

pub use driver::{DevicePtr, DriverError, HipContext, HipFunction, HipModule, HipSlice, HipStream};
pub use hiprtc::{HiprtcError, Hsaco};

#[cfg(feature = "rocblas")]
pub use rocblas::{
    Axpy, AxpyConfig, Copy, CopyConfig, Dot, DotConfig, Gemm, GemmConfig, Gemv, GemvConfig, Nrm2,
    Nrm2Config, Operation, RocblasError, RocblasHandle, Scal, ScalConfig, StridedBatchedConfig,
};

#[cfg(feature = "hipblaslt")]
pub use hipblaslt::{
    HipBlasLt, HipBlasLtError, MatmulDesc, MatmulHeuristic, MatmulPref, MatrixLayout,
};

#[cfg(feature = "dynamic-loading")]
pub fn get_lib_name_candidates(lib_name: &str) -> Vec<String> {
    use std::env::consts::{DLL_PREFIX, DLL_SUFFIX};

    let major = env!("ROCM_MAJOR_VERSION");
    let minor = env!("ROCM_MINOR_VERSION");

    let rocm_path =
        std::env::var("ROCM_PATH").unwrap_or_else(|_| "/opt/rocm".to_string());

    let stems = [
        // Bare soname; resolved via LD_LIBRARY_PATH / ldconfig.
        format!("{DLL_PREFIX}{lib_name}{DLL_SUFFIX}"),
        // Major-versioned soname (e.g. librocblas.so.4).
        format!("{DLL_PREFIX}{lib_name}{DLL_SUFFIX}.{major}"),
        // Common alternative majors to try as fallbacks.
        format!("{DLL_PREFIX}{lib_name}{DLL_SUFFIX}.0"),
        format!("{DLL_PREFIX}{lib_name}{DLL_SUFFIX}.1"),
        // Exact patch soname.
        format!("{DLL_PREFIX}{lib_name}{DLL_SUFFIX}.{major}.{minor}"),
    ];

    let mut out = Vec::with_capacity(stems.len() * 2);
    for stem in &stems {
        out.push(stem.clone());                    // search via loader
        out.push(format!("{rocm_path}/lib/{stem}")); // explicit ROCm install
    }
    out
}

#[cfg(feature = "dynamic-loading")]
pub fn panic_no_lib_found(name: &str, tried: &[String]) -> ! {
    panic!(
        "could not load dynamic library {name}; tried {} candidates: {tried:?}",
        tried.len()
    );
}