rocmrc 0.4.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
66
67
68
69
70
71

//! 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 hip;
pub mod hiprtc;

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

pub use hip::{
    DevicePtr, DevicePtrMut, DeviceRepr, DeviceSlice, EventWaitFlags, ExternalMemory, HipContext,
    HipError, HipEvent, HipFunction, HipGraph, HipModule, HipSlice, HipStream, HipUnifiedSlice,
    HipUnifiedView, HipUnifiedViewMut, HipView, HipViewMut, HostSlice, LaunchArgs, LaunchConfig,
    MappedBuffer, MemAttachFlags, PinnedHostSlice, Profiler, PushKernelArg, StreamKind, SyncOnDrop,
    ValidAsZeroBits, profiler_start, profiler_stop,
};
pub use hiprtc::{HiprtcError, Hsaco};

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

#[cfg(feature = "hipblaslt")]
pub use hipblaslt::{
    Activation, HipBlasLT, HipblasError, Matmul, MatmulConfig, MatmulShared, Workspace,
};

#[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()
    );
}