sim_lib_compute_rocm/
loader.rs

1//! Runtime ROCm/rocBLAS symbol discovery.
2
3use std::{
4    ffi::c_void,
5    fmt,
6    path::{Path, PathBuf},
7    process::Command,
8    sync::Arc,
9};
10
11use libloading::Library;
12
13const HIP_NAMES: &[&str] = &["libamdhip64.so.7", "libamdhip64.so.6", "libamdhip64.so"];
14const ROCBLAS_NAMES: &[&str] = &[
15    "librocblas.so.5",
16    "librocblas.so.4",
17    "librocblas.so.0",
18    "librocblas.so",
19];
20const ROCBLASLT_NAMES: &[&str] = &["librocblaslt.so.0", "librocblaslt.so"];
21
22const HIP_SYMBOLS: &[&str] = &[
23    "hipInit",
24    "hipRuntimeGetVersion",
25    "hipGetDeviceCount",
26    "hipMalloc",
27    "hipFree",
28    "hipMemcpy",
29    "hipDeviceSynchronize",
30];
31const ROCBLAS_SYMBOLS: &[&str] = &[
32    "rocblas_create_handle",
33    "rocblas_destroy_handle",
34    "rocblas_sgemm",
35    "rocblas_gemm_ex",
36];
37const ROCBLASLT_SYMBOLS: &[&str] = &["rocblaslt_create_handle", "rocblaslt_destroy_handle"];
38
39/// One validated runtime symbol.
40#[derive(Clone, Debug, PartialEq, Eq)]
41pub struct RocmSymbolEvidence {
42    /// Symbol name.
43    pub name: String,
44    /// Whether the dynamic library exported the symbol.
45    pub present: bool,
46}
47
48/// Dynamic-library ABI evidence required by the ROCm provider.
49#[derive(Clone, Debug, PartialEq, Eq)]
50pub struct RocmAbiEvidence {
51    /// Loaded HIP runtime library path or platform name.
52    pub hip_library: String,
53    /// Loaded rocBLAS library path or platform name.
54    pub rocblas_library: String,
55    /// Loaded rocBLASLt library path or platform name, when present.
56    pub rocblaslt_library: Option<String>,
57    /// HIP runtime version when the runtime can report it.
58    pub hip_runtime_version: Option<i32>,
59    /// Observed AMD GPU ISA targets such as `gfx1103`.
60    pub observed_gfx_targets: Vec<String>,
61    /// Checked HIP runtime symbols.
62    pub hip_symbols: Vec<RocmSymbolEvidence>,
63    /// Checked rocBLAS symbols.
64    pub rocblas_symbols: Vec<RocmSymbolEvidence>,
65    /// Checked rocBLASLt symbols.
66    pub rocblaslt_symbols: Vec<RocmSymbolEvidence>,
67}
68
69impl RocmAbiEvidence {
70    /// Returns true when Linux, HIP, rocBLAS, and a concrete gfx target exist.
71    pub fn is_complete(&self) -> bool {
72        cfg!(target_os = "linux")
73            && !self.observed_gfx_targets.is_empty()
74            && self.hip_symbols.iter().all(|symbol| symbol.present)
75            && self.rocblas_symbols.iter().all(|symbol| symbol.present)
76    }
77
78    /// Returns true when half-family matmul may use the validated rocBLASLt path.
79    pub fn supports_half_matmul(&self) -> bool {
80        self.rocblas_symbols
81            .iter()
82            .any(|symbol| symbol.name == "rocblas_gemm_ex" && symbol.present)
83            && !self.rocblaslt_symbols.is_empty()
84            && self.rocblaslt_symbols.iter().all(|symbol| symbol.present)
85    }
86}
87
88/// Loaded ROCm runtime libraries kept alive for function-pointer validity.
89pub struct RocmLibrarySet {
90    evidence: RocmAbiEvidence,
91    hip: Library,
92    rocblas: Library,
93    rocblaslt: Option<Library>,
94}
95
96impl RocmLibrarySet {
97    fn new(
98        evidence: RocmAbiEvidence,
99        hip: Library,
100        rocblas: Library,
101        rocblaslt: Option<Library>,
102    ) -> Self {
103        Self {
104            evidence,
105            hip,
106            rocblas,
107            rocblaslt,
108        }
109    }
110
111    /// Returns checked ABI evidence.
112    pub fn evidence(&self) -> &RocmAbiEvidence {
113        &self.evidence
114    }
115
116    /// Returns loaded library handles to keep symbols alive.
117    pub fn handles(&self) -> (&Library, &Library, Option<&Library>) {
118        (&self.hip, &self.rocblas, self.rocblaslt.as_ref())
119    }
120
121    pub(crate) fn execution_handles(&self) -> (&Library, &Library) {
122        (&self.hip, &self.rocblas)
123    }
124}
125
126impl fmt::Debug for RocmLibrarySet {
127    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
128        formatter
129            .debug_struct("RocmLibrarySet")
130            .field("evidence", &self.evidence)
131            .finish_non_exhaustive()
132    }
133}
134
135/// Result of ROCm runtime discovery.
136#[derive(Clone, Debug)]
137pub struct RocmRuntimeProbe {
138    /// Validated loaded runtime, when discovery succeeded.
139    pub runtime: Option<Arc<RocmLibrarySet>>,
140    /// ABI evidence from the successful runtime or the best failed probe.
141    pub evidence: Option<RocmAbiEvidence>,
142    /// Diagnostics collected while searching dynamic libraries.
143    pub diagnostics: Vec<String>,
144}
145
146impl RocmRuntimeProbe {
147    /// Builds a successful probe from validated evidence without library
148    /// handles. This is intended for deterministic fake-loader tests.
149    pub fn fake_present(evidence: RocmAbiEvidence) -> Self {
150        Self {
151            runtime: None,
152            evidence: Some(evidence),
153            diagnostics: Vec::new(),
154        }
155    }
156
157    /// Returns true when discovery validated a usable ROCm provider.
158    pub fn is_available(&self) -> bool {
159        self.evidence
160            .as_ref()
161            .is_some_and(RocmAbiEvidence::is_complete)
162    }
163}
164
165/// ROCm dynamic-loading failure.
166#[derive(Clone, Debug, PartialEq, Eq)]
167pub struct RocmLoadError {
168    /// Human-readable failure message.
169    pub message: String,
170}
171
172impl fmt::Display for RocmLoadError {
173    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
174        formatter.write_str(&self.message)
175    }
176}
177
178impl std::error::Error for RocmLoadError {}
179
180/// Loader abstraction used by real and fake ROCm discovery.
181pub trait DynamicRocmLoader {
182    /// Performs ROCm runtime discovery.
183    fn discover(&self) -> Result<RocmRuntimeProbe, RocmLoadError>;
184}
185
186/// Real dynamic loader using platform ROCm shared libraries.
187#[derive(Clone, Debug)]
188pub struct RocmRuntimeLoader {
189    search_dirs: Vec<PathBuf>,
190    search_system: bool,
191}
192
193impl Default for RocmRuntimeLoader {
194    fn default() -> Self {
195        Self {
196            search_dirs: Vec::new(),
197            search_system: true,
198        }
199    }
200}
201
202impl RocmRuntimeLoader {
203    /// Builds a loader that searches platform library paths.
204    pub fn new() -> Self {
205        Self::default()
206    }
207
208    /// Builds a loader that first searches explicit directories.
209    pub fn with_search_dirs(search_dirs: Vec<PathBuf>) -> Self {
210        Self {
211            search_dirs,
212            search_system: true,
213        }
214    }
215
216    /// Builds a loader restricted to explicit directories. This is used to
217    /// prove fail-closed behavior when vendor libraries are unavailable.
218    pub fn with_search_dirs_only(search_dirs: Vec<PathBuf>) -> Self {
219        Self {
220            search_dirs,
221            search_system: false,
222        }
223    }
224}
225
226impl DynamicRocmLoader for RocmRuntimeLoader {
227    fn discover(&self) -> Result<RocmRuntimeProbe, RocmLoadError> {
228        let mut diagnostics = Vec::new();
229        if !cfg!(target_os = "linux") {
230            return Err(RocmLoadError {
231                message: "ROCm provider is supported only on Linux".to_owned(),
232            });
233        }
234        let (hip_name, hip) = self.open_first(HIP_NAMES, &mut diagnostics)?;
235        let (rocblas_name, rocblas) = self.open_first(ROCBLAS_NAMES, &mut diagnostics)?;
236        let rocblaslt = self.open_first(ROCBLASLT_NAMES, &mut diagnostics).ok();
237
238        let hip_symbols = symbol_evidence(&hip, HIP_SYMBOLS);
239        let rocblas_symbols = symbol_evidence(&rocblas, ROCBLAS_SYMBOLS);
240        let rocblaslt_symbols = rocblaslt
241            .as_ref()
242            .map(|(_, library)| symbol_evidence(library, ROCBLASLT_SYMBOLS))
243            .unwrap_or_default();
244        let hip_runtime_version = hip_runtime_version(&hip).ok();
245        let observed_gfx_targets = observed_gfx_targets();
246        let evidence = RocmAbiEvidence {
247            hip_library: hip_name,
248            rocblas_library: rocblas_name,
249            rocblaslt_library: rocblaslt.as_ref().map(|(name, _)| name.clone()),
250            hip_runtime_version,
251            observed_gfx_targets,
252            hip_symbols,
253            rocblas_symbols,
254            rocblaslt_symbols,
255        };
256        if !evidence.is_complete() {
257            return Ok(RocmRuntimeProbe {
258                runtime: None,
259                evidence: Some(evidence),
260                diagnostics,
261            });
262        }
263        let runtime = Arc::new(RocmLibrarySet::new(
264            evidence.clone(),
265            hip,
266            rocblas,
267            rocblaslt.map(|(_, library)| library),
268        ));
269        Ok(RocmRuntimeProbe {
270            runtime: Some(runtime),
271            evidence: Some(evidence),
272            diagnostics,
273        })
274    }
275}
276
277impl RocmRuntimeLoader {
278    fn open_first(
279        &self,
280        names: &[&str],
281        diagnostics: &mut Vec<String>,
282    ) -> Result<(String, Library), RocmLoadError> {
283        for name in candidate_paths(&self.search_dirs, names, self.search_system) {
284            match open_library(&name) {
285                Ok(library) => return Ok((name.display().to_string(), library)),
286                Err(error) => diagnostics.push(format!("{}: {error}", name.display())),
287            }
288        }
289        Err(RocmLoadError {
290            message: format!("ROCm library was not found; tried {}", names.join(", ")),
291        })
292    }
293}
294
295/// Fake loader for deterministic tests.
296#[derive(Clone, Debug)]
297pub struct FakeRocmLoader {
298    probe: Result<RocmRuntimeProbe, RocmLoadError>,
299}
300
301impl FakeRocmLoader {
302    /// Builds a fake loader that returns validated ROCm evidence.
303    pub fn available() -> Self {
304        Self {
305            probe: Ok(RocmRuntimeProbe::fake_present(complete_fake_evidence())),
306        }
307    }
308
309    /// Builds a fake loader with incomplete core HIP/rocBLAS ABI evidence.
310    pub fn incomplete() -> Self {
311        let mut evidence = complete_fake_evidence();
312        if let Some(symbol) = evidence
313            .rocblas_symbols
314            .iter_mut()
315            .find(|symbol| symbol.name == "rocblas_sgemm")
316        {
317            symbol.present = false;
318        }
319        Self {
320            probe: Ok(RocmRuntimeProbe {
321                runtime: None,
322                evidence: Some(evidence),
323                diagnostics: vec!["missing rocblas_sgemm".to_owned()],
324            }),
325        }
326    }
327
328    /// Builds a fake loader with missing optional rocBLASLt evidence.
329    pub fn without_rocblaslt() -> Self {
330        let mut evidence = complete_fake_evidence();
331        if let Some(symbol) = evidence
332            .rocblaslt_symbols
333            .iter_mut()
334            .find(|symbol| symbol.name == "rocblaslt_create_handle")
335        {
336            symbol.present = false;
337        }
338        Self {
339            probe: Ok(RocmRuntimeProbe {
340                runtime: None,
341                evidence: Some(evidence),
342                diagnostics: vec!["missing rocblaslt_create_handle".to_owned()],
343            }),
344        }
345    }
346
347    /// Builds a fake loader that reports ROCm as absent.
348    pub fn absent() -> Self {
349        Self {
350            probe: Err(RocmLoadError {
351                message: "ROCm runtime absent".to_owned(),
352            }),
353        }
354    }
355}
356
357impl DynamicRocmLoader for FakeRocmLoader {
358    fn discover(&self) -> Result<RocmRuntimeProbe, RocmLoadError> {
359        self.probe.clone()
360    }
361}
362
363/// Discovers ROCm using the real platform dynamic loader.
364pub fn discover_rocm_runtime() -> Result<RocmRuntimeProbe, RocmLoadError> {
365    RocmRuntimeLoader::new().discover()
366}
367
368fn complete_fake_evidence() -> RocmAbiEvidence {
369    RocmAbiEvidence {
370        hip_library: "fake-libamdhip64".to_owned(),
371        rocblas_library: "fake-librocblas".to_owned(),
372        rocblaslt_library: Some("fake-librocblaslt".to_owned()),
373        hip_runtime_version: Some(6_300_000),
374        observed_gfx_targets: vec!["gfx1103".to_owned()],
375        hip_symbols: HIP_SYMBOLS
376            .iter()
377            .map(|name| RocmSymbolEvidence {
378                name: (*name).to_owned(),
379                present: true,
380            })
381            .collect(),
382        rocblas_symbols: ROCBLAS_SYMBOLS
383            .iter()
384            .map(|name| RocmSymbolEvidence {
385                name: (*name).to_owned(),
386                present: true,
387            })
388            .collect(),
389        rocblaslt_symbols: ROCBLASLT_SYMBOLS
390            .iter()
391            .map(|name| RocmSymbolEvidence {
392                name: (*name).to_owned(),
393                present: true,
394            })
395            .collect(),
396    }
397}
398
399fn candidate_paths(search_dirs: &[PathBuf], names: &[&str], search_system: bool) -> Vec<PathBuf> {
400    let mut candidates = Vec::new();
401    for directory in search_dirs {
402        for name in names {
403            candidates.push(directory.join(name));
404        }
405    }
406    if search_system {
407        candidates.extend(names.iter().map(PathBuf::from));
408    }
409    candidates
410}
411
412fn symbol_evidence(library: &Library, names: &[&str]) -> Vec<RocmSymbolEvidence> {
413    names
414        .iter()
415        .map(|name| RocmSymbolEvidence {
416            name: (*name).to_owned(),
417            present: symbol_present(library, name),
418        })
419        .collect()
420}
421
422fn open_library(path: &Path) -> Result<Library, libloading::Error> {
423    // SAFETY: The handle remains owned by RocmLibrarySet while symbols are used.
424    unsafe { Library::new(path) }
425}
426
427fn symbol_present(library: &Library, name: &str) -> bool {
428    let mut bytes = name.as_bytes().to_vec();
429    bytes.push(0);
430    // SAFETY: The lookup only checks whether the library exports the named
431    // symbol as an opaque address. The address is not called or dereferenced.
432    unsafe { library.get::<*mut c_void>(&bytes).is_ok() }
433}
434
435fn hip_runtime_version(library: &Library) -> Result<i32, RocmLoadError> {
436    type HipInit = unsafe extern "C" fn(u32) -> i32;
437    type HipRuntimeGetVersion = unsafe extern "C" fn(*mut i32) -> i32;
438    type HipGetDeviceCount = unsafe extern "C" fn(*mut i32) -> i32;
439    // SAFETY: These are the documented HIP signatures. Pointers are initialized,
440    // and only status 0 is accepted.
441    unsafe {
442        let hip_init = library
443            .get::<HipInit>(b"hipInit\0")
444            .map_err(|error| RocmLoadError {
445                message: error.to_string(),
446            })?;
447        let get_version = library
448            .get::<HipRuntimeGetVersion>(b"hipRuntimeGetVersion\0")
449            .map_err(|error| RocmLoadError {
450                message: error.to_string(),
451            })?;
452        let get_device_count = library
453            .get::<HipGetDeviceCount>(b"hipGetDeviceCount\0")
454            .map_err(|error| RocmLoadError {
455                message: error.to_string(),
456            })?;
457        let init_status = hip_init(0);
458        if init_status != 0 {
459            return Err(RocmLoadError {
460                message: format!("hipInit failed with status {init_status}"),
461            });
462        }
463        let mut device_count = 0;
464        let count_status = get_device_count(&mut device_count);
465        if count_status != 0 || device_count <= 0 {
466            return Err(RocmLoadError {
467                message: format!("hipGetDeviceCount failed with status {count_status}"),
468            });
469        }
470        let mut version = 0;
471        let version_status = get_version(&mut version);
472        if version_status != 0 {
473            return Err(RocmLoadError {
474                message: format!("hipRuntimeGetVersion failed with status {version_status}"),
475            });
476        }
477        Ok(version)
478    }
479}
480
481fn observed_gfx_targets() -> Vec<String> {
482    let Ok(output) = Command::new("rocm_agent_enumerator").output() else {
483        return Vec::new();
484    };
485    if !output.status.success() {
486        return Vec::new();
487    }
488    let Ok(stdout) = String::from_utf8(output.stdout) else {
489        return Vec::new();
490    };
491    let mut targets = stdout
492        .lines()
493        .map(str::trim)
494        .filter(|line| line.starts_with("gfx") && *line != "gfx000")
495        .map(ToOwned::to_owned)
496        .collect::<Vec<_>>();
497    targets.sort();
498    targets.dedup();
499    targets
500}
sim_lib_compute_rocm/loader.rs

sim_lib_compute_rocm/
loader.rs
