sim_lib_compute_rocm/
site.rs

1//! Loadable ROCm compute site exports.
2
3use std::sync::{Arc, Mutex};
4
5use sim_kernel::{
6    AbiVersion, CapabilityName, DefaultFactory, Export, Factory, Lib, LibManifest, LibTarget,
7    Linker, Result, Symbol, Version,
8};
9use sim_lib_numbers_tensor::{
10    CpuTensorExecutor, SubmissionEvidence, Tensor, TensorExecError, TensorExecution,
11    TensorExecutor, TensorExecutorCard, TensorRequest, TensorSite, domains, matmul_exec_op_symbol,
12    parse_f32_literal_cell,
13};
14
15use crate::{
16    DynamicRocmLoader, RocmAbiEvidence, RocmAllocation, RocmLibrarySet, RocmLoadError,
17    RocmResidentStorage, RocmRuntimeProbe, discover_rocm_runtime, runtime::RocmDeviceBuffer,
18};
19
20/// Stable symbol for the ROCm runtime library.
21pub fn compute_rocm_lib_symbol() -> Symbol {
22    Symbol::qualified("compute", "rocm-lib")
23}
24
25/// Stable symbol for the ROCm tensor executor.
26pub fn rocm_executor_symbol() -> Symbol {
27    Symbol::qualified("compute", "executor/rocm")
28}
29
30/// Site symbol exported by a successful ROCm runtime.
31pub fn compute_rocm_site_symbol() -> Symbol {
32    Symbol::new("site/compute/rocm")
33}
34
35/// Capability required before a ROCm hardware tensor site can be realized.
36pub fn compute_rocm_capability() -> CapabilityName {
37    CapabilityName::new("device.gpu.rocm")
38}
39
40#[derive(Clone, Debug, Default, PartialEq, Eq)]
41struct RocmExecutorState {
42    accepted: usize,
43    queued: usize,
44    next_allocation: usize,
45}
46
47/// Tensor executor backed by validated ROCm/rocBLAS ABI evidence.
48#[derive(Clone)]
49pub struct RocmTensorExecutor {
50    evidence: RocmAbiEvidence,
51    runtime: Option<Arc<RocmLibrarySet>>,
52    state: Arc<Mutex<RocmExecutorState>>,
53}
54
55impl RocmTensorExecutor {
56    /// Builds an executor from validated ROCm ABI evidence.
57    pub fn new(evidence: RocmAbiEvidence) -> Self {
58        Self {
59            evidence,
60            runtime: None,
61            state: Arc::new(Mutex::new(RocmExecutorState::default())),
62        }
63    }
64
65    /// Builds an executor that submits f32 matmul to a validated ROCm runtime.
66    pub fn from_runtime(runtime: Arc<RocmLibrarySet>) -> Self {
67        Self {
68            evidence: runtime.evidence().clone(),
69            runtime: Some(runtime),
70            state: Arc::new(Mutex::new(RocmExecutorState::default())),
71        }
72    }
73
74    /// Returns the ROCm ABI evidence used by this executor.
75    pub fn evidence(&self) -> &RocmAbiEvidence {
76        &self.evidence
77    }
78
79    fn dtype_supported(&self, dtype: &Symbol) -> bool {
80        dtype == &domains::f32()
81    }
82
83    fn reserve_allocation(
84        &self,
85        shape: &[usize],
86        operation: Symbol,
87    ) -> std::result::Result<RocmAllocation, TensorExecError> {
88        let bytes = tensor_bytes(shape)?;
89        let mut state = self.state.lock().expect("rocm executor state poisoned");
90        state.accepted += 1;
91        state.queued += 1;
92        state.next_allocation += 1;
93        Ok(RocmAllocation {
94            id: state.next_allocation,
95            bytes,
96            operation,
97        })
98    }
99
100    fn execute_runtime(
101        &self,
102        runtime: &Arc<RocmLibrarySet>,
103        request: &TensorRequest,
104        allocation: RocmAllocation,
105    ) -> std::result::Result<TensorExecution, TensorExecError> {
106        let [left, right] = request.inputs.as_ref() else {
107            return Err(invalid("rocm matmul requires two inputs"));
108        };
109        let [rows, inner] = left.shape() else {
110            return Err(invalid("rocm matmul left input must be rank two"));
111        };
112        let [right_inner, cols] = right.shape() else {
113            return Err(invalid("rocm matmul right input must be rank two"));
114        };
115        if inner != right_inner || request.output.shape() != [*rows, *cols] {
116            return Err(invalid("rocm matmul shapes do not conform"));
117        }
118        let left = rocm_input(runtime, left)?;
119        let right = rocm_input(runtime, right)?;
120        let output = runtime
121            .matmul(&left, &right, *rows, *inner, *cols)
122            .map_err(execution_error)?;
123        let storage = RocmResidentStorage::from_device(
124            compute_rocm_site_symbol(),
125            allocation,
126            request.output.shape().to_vec(),
127            request.output.dtype().clone(),
128            output,
129        );
130        Ok(TensorExecution::Complete(Tensor::from_storage(
131            request.output.shape().to_vec(),
132            request.output.dtype().clone(),
133            Arc::new(storage),
134        )?))
135    }
136}
137
138impl TensorExecutor for RocmTensorExecutor {
139    fn card(&self) -> TensorExecutorCard {
140        TensorExecutorCard::new(
141            rocm_executor_symbol(),
142            "rocm/rocblas",
143            Symbol::qualified("compute", "rocm"),
144            vec![matmul_exec_op_symbol()],
145            Some(compute_rocm_capability()),
146        )
147    }
148
149    fn execute(
150        &self,
151        cx: &mut sim_kernel::Cx,
152        request: TensorRequest,
153    ) -> std::result::Result<TensorExecution, TensorExecError> {
154        if request.operation.symbol != matmul_exec_op_symbol() {
155            return Ok(TensorExecution::Unsupported {
156                reason: Arc::from("rocm provider accepts dense matmul only"),
157            });
158        }
159        if !self.dtype_supported(request.output.dtype()) {
160            return Ok(TensorExecution::Unsupported {
161                reason: Arc::from("rocm provider accepts dense f32 matmul"),
162            });
163        }
164        let allocation =
165            self.reserve_allocation(request.output.shape(), request.operation.symbol.clone())?;
166        if let Some(runtime) = &self.runtime {
167            return self.execute_runtime(runtime, &request, allocation);
168        }
169        let result = CpuTensorExecutor::new().execute(cx, request)?;
170        let TensorExecution::Complete(tensor) = result else {
171            return Ok(result);
172        };
173        resident_result(tensor, allocation)
174    }
175
176    fn flush(&self) -> std::result::Result<SubmissionEvidence, TensorExecError> {
177        let mut state = self.state.lock().expect("rocm executor state poisoned");
178        let accepted = state.queued;
179        state.queued = 0;
180        Ok(SubmissionEvidence::new(rocm_executor_symbol(), accepted))
181    }
182}
183
184fn resident_result(
185    tensor: Tensor,
186    allocation: RocmAllocation,
187) -> std::result::Result<TensorExecution, TensorExecError> {
188    let cells = tensor.cells().map_err(TensorExecError::from)?;
189    let storage = RocmResidentStorage::new(
190        compute_rocm_site_symbol(),
191        allocation,
192        tensor.shape().to_vec(),
193        tensor.dtype().clone(),
194        cells,
195    );
196    Ok(TensorExecution::Complete(Tensor::from_storage(
197        tensor.shape().to_vec(),
198        tensor.dtype().clone(),
199        Arc::new(storage),
200    )?))
201}
202
203fn tensor_bytes(shape: &[usize]) -> std::result::Result<u64, TensorExecError> {
204    let cells = shape.iter().try_fold(1_u64, |count, extent| {
205        count
206            .checked_mul(u64::try_from(*extent).map_err(|_| invalid("rocm extent exceeds u64"))?)
207            .ok_or_else(|| invalid("rocm tensor byte count overflowed"))
208    })?;
209    cells
210        .checked_mul(4)
211        .ok_or_else(|| invalid("rocm tensor byte count overflowed"))
212}
213
214fn invalid(message: impl Into<Arc<str>>) -> TensorExecError {
215    TensorExecError::InvalidRequest {
216        message: message.into(),
217    }
218}
219
220/// Loadable library that registers a ROCm compute site when runtime validation
221/// succeeds.
222#[derive(Clone, Debug, Default)]
223pub struct ComputeRocmLib {
224    probe: Option<RocmRuntimeProbe>,
225}
226
227impl ComputeRocmLib {
228    /// Probes local ROCm dynamic libraries and builds a provider library.
229    pub fn probe() -> std::result::Result<Self, RocmLoadError> {
230        Ok(Self {
231            probe: Some(discover_rocm_runtime()?),
232        })
233    }
234
235    /// Builds a provider from an injected loader.
236    pub fn from_loader(loader: &dyn DynamicRocmLoader) -> std::result::Result<Self, RocmLoadError> {
237        Ok(Self {
238            probe: Some(loader.discover()?),
239        })
240    }
241
242    /// Builds a provider from precomputed probe evidence.
243    pub fn from_probe(probe: RocmRuntimeProbe) -> Self {
244        Self { probe: Some(probe) }
245    }
246
247    /// Returns ROCm runtime probe evidence.
248    pub fn probe_evidence(&self) -> Option<&RocmRuntimeProbe> {
249        self.probe.as_ref()
250    }
251
252    fn available_runtime(&self) -> Option<Arc<RocmLibrarySet>> {
253        self.probe
254            .as_ref()
255            .and_then(|probe| probe.runtime.clone())
256            .filter(|runtime| runtime.evidence().is_complete())
257    }
258}
259
260impl Lib for ComputeRocmLib {
261    fn manifest(&self) -> LibManifest {
262        LibManifest {
263            id: compute_rocm_lib_symbol(),
264            version: Version(env!("CARGO_PKG_VERSION").to_owned()),
265            abi: AbiVersion { major: 0, minor: 1 },
266            target: LibTarget::HostRegistered,
267            requires: Vec::new(),
268            capabilities: self
269                .available_runtime()
270                .map(|_| vec![compute_rocm_capability()])
271                .unwrap_or_default(),
272            exports: self
273                .available_runtime()
274                .map(|_| {
275                    vec![Export::Site {
276                        symbol: compute_rocm_site_symbol(),
277                        runtime_id: None,
278                    }]
279                })
280                .unwrap_or_default(),
281        }
282    }
283
284    fn load(&self, _cx: &mut sim_kernel::LoadCx, linker: &mut Linker<'_>) -> Result<()> {
285        let Some(runtime) = self.available_runtime() else {
286            return Ok(());
287        };
288        let executor = Arc::new(RocmTensorExecutor::from_runtime(runtime));
289        let site = TensorSite::new(
290            compute_rocm_site_symbol(),
291            executor,
292            vec![compute_rocm_capability()],
293        );
294        linker.site_value(
295            compute_rocm_site_symbol(),
296            DefaultFactory.opaque(Arc::new(site))?,
297        )?;
298        Ok(())
299    }
300}
301
302fn rocm_input(
303    runtime: &Arc<RocmLibrarySet>,
304    tensor: &Tensor,
305) -> std::result::Result<Arc<RocmDeviceBuffer>, TensorExecError> {
306    if let Some(buffer) = tensor
307        .storage()
308        .as_any()
309        .downcast_ref::<RocmResidentStorage>()
310        .and_then(RocmResidentStorage::device_buffer)
311        .filter(|buffer| Arc::ptr_eq(buffer.runtime(), runtime))
312    {
313        return Ok(Arc::clone(buffer));
314    }
315    let values = tensor
316        .cells()
317        .map_err(TensorExecError::from)?
318        .iter()
319        .map(|cell| {
320            parse_f32_literal_cell(cell)
321                .ok_or_else(|| invalid("rocm matmul input is not canonical f32"))
322        })
323        .collect::<std::result::Result<Vec<_>, _>>()?;
324    runtime.upload(&values).map_err(execution_error)
325}
326
327fn execution_error(error: RocmLoadError) -> TensorExecError {
328    TensorExecError::Eval {
329        message: Arc::from(error.to_string()),
330    }
331}
sim_lib_compute_rocm/site.rs

sim_lib_compute_rocm/
site.rs
