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

sim_lib_compute_rocm/
loader.rs
