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

sim_lib_compute_rocm/
loader.rs
