sim_lib_compute_cuda/
site.rs

1//! Loadable CUDA 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    CudaAbiEvidence, CudaAllocation, CudaLibrarySet, CudaLoadError, CudaResidentStorage,
17    CudaRuntimeProbe, DynamicCudaLoader, discover_cuda_runtime, runtime::CudaDeviceBuffer,
18};
19
20/// Stable symbol for the CUDA runtime library.
21pub fn compute_cuda_lib_symbol() -> Symbol {
22    Symbol::qualified("compute", "cuda-lib")
23}
24
25/// Stable symbol for the CUDA tensor executor.
26pub fn cuda_executor_symbol() -> Symbol {
27    Symbol::qualified("compute", "executor/cuda")
28}
29
30/// Site symbol exported by a successful CUDA runtime.
31pub fn compute_cuda_site_symbol() -> Symbol {
32    Symbol::new("site/compute/cuda")
33}
34
35/// Capability required before a CUDA hardware tensor site can be realized.
36pub fn compute_cuda_capability() -> CapabilityName {
37    CapabilityName::new("device.gpu.cuda")
38}
39
40#[derive(Clone, Debug, Default, PartialEq, Eq)]
41struct CudaExecutorState {
42    accepted: usize,
43    queued: usize,
44    next_allocation: usize,
45}
46
47/// Tensor executor backed by validated CUDA/cuBLAS ABI evidence.
48#[derive(Clone)]
49pub struct CudaTensorExecutor {
50    evidence: CudaAbiEvidence,
51    runtime: Option<Arc<CudaLibrarySet>>,
52    state: Arc<Mutex<CudaExecutorState>>,
53}
54
55impl CudaTensorExecutor {
56    /// Builds an executor from validated CUDA ABI evidence.
57    pub fn new(evidence: CudaAbiEvidence) -> Self {
58        Self {
59            evidence,
60            runtime: None,
61            state: Arc::new(Mutex::new(CudaExecutorState::default())),
62        }
63    }
64
65    /// Builds an executor that submits f32 matmul to a validated CUDA runtime.
66    pub fn from_runtime(runtime: Arc<CudaLibrarySet>) -> Self {
67        Self {
68            evidence: runtime.evidence().clone(),
69            runtime: Some(runtime),
70            state: Arc::new(Mutex::new(CudaExecutorState::default())),
71        }
72    }
73
74    /// Returns the CUDA ABI evidence used by this executor.
75    pub fn evidence(&self) -> &CudaAbiEvidence {
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<CudaAllocation, TensorExecError> {
88        let bytes = tensor_bytes(shape)?;
89        let mut state = self.state.lock().expect("cuda executor state poisoned");
90        state.accepted += 1;
91        state.queued += 1;
92        state.next_allocation += 1;
93        Ok(CudaAllocation {
94            id: state.next_allocation,
95            bytes,
96            operation,
97        })
98    }
99
100    fn execute_runtime(
101        &self,
102        runtime: &Arc<CudaLibrarySet>,
103        request: &TensorRequest,
104        allocation: CudaAllocation,
105    ) -> std::result::Result<TensorExecution, TensorExecError> {
106        let [left, right] = request.inputs.as_ref() else {
107            return Err(invalid("cuda matmul requires two inputs"));
108        };
109        let [rows, inner] = left.shape() else {
110            return Err(invalid("cuda matmul left input must be rank two"));
111        };
112        let [right_inner, cols] = right.shape() else {
113            return Err(invalid("cuda matmul right input must be rank two"));
114        };
115        if inner != right_inner || request.output.shape() != [*rows, *cols] {
116            return Err(invalid("cuda matmul shapes do not conform"));
117        }
118        let left = cuda_input(runtime, left)?;
119        let right = cuda_input(runtime, right)?;
120        let output = runtime
121            .matmul(&left, &right, *rows, *inner, *cols)
122            .map_err(execution_error)?;
123        let storage = CudaResidentStorage::from_device(
124            compute_cuda_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 CudaTensorExecutor {
139    fn card(&self) -> TensorExecutorCard {
140        TensorExecutorCard::new(
141            cuda_executor_symbol(),
142            "cuda/cublas",
143            Symbol::qualified("compute", "cuda"),
144            vec![matmul_exec_op_symbol()],
145            Some(compute_cuda_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("cuda provider accepts dense matmul only"),
157            });
158        }
159        if !self.dtype_supported(request.output.dtype()) {
160            return Ok(TensorExecution::Unsupported {
161                reason: Arc::from("cuda 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("cuda executor state poisoned");
178        let accepted = state.queued;
179        state.queued = 0;
180        Ok(SubmissionEvidence::new(cuda_executor_symbol(), accepted))
181    }
182}
183
184fn resident_result(
185    tensor: Tensor,
186    allocation: CudaAllocation,
187) -> std::result::Result<TensorExecution, TensorExecError> {
188    let cells = tensor.cells().map_err(TensorExecError::from)?;
189    let storage = CudaResidentStorage::new(
190        compute_cuda_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("cuda extent exceeds u64"))?)
207            .ok_or_else(|| invalid("cuda tensor byte count overflowed"))
208    })?;
209    cells
210        .checked_mul(4)
211        .ok_or_else(|| invalid("cuda 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 CUDA compute site when runtime validation
221/// succeeds.
222#[derive(Clone, Debug, Default)]
223pub struct ComputeCudaLib {
224    probe: Option<CudaRuntimeProbe>,
225}
226
227impl ComputeCudaLib {
228    /// Probes local CUDA dynamic libraries and builds a provider library.
229    pub fn probe() -> std::result::Result<Self, CudaLoadError> {
230        Ok(Self {
231            probe: Some(discover_cuda_runtime()?),
232        })
233    }
234
235    /// Builds a provider from an injected loader.
236    pub fn from_loader(loader: &dyn DynamicCudaLoader) -> std::result::Result<Self, CudaLoadError> {
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: CudaRuntimeProbe) -> Self {
244        Self { probe: Some(probe) }
245    }
246
247    /// Returns CUDA runtime probe evidence.
248    pub fn probe_evidence(&self) -> Option<&CudaRuntimeProbe> {
249        self.probe.as_ref()
250    }
251
252    fn available_runtime(&self) -> Option<Arc<CudaLibrarySet>> {
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 ComputeCudaLib {
261    fn manifest(&self) -> LibManifest {
262        LibManifest {
263            id: compute_cuda_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_cuda_capability()])
271                .unwrap_or_default(),
272            exports: self
273                .available_runtime()
274                .map(|_| {
275                    vec![Export::Site {
276                        symbol: compute_cuda_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(CudaTensorExecutor::from_runtime(runtime));
289        let site = TensorSite::new(
290            compute_cuda_site_symbol(),
291            executor,
292            vec![compute_cuda_capability()],
293        );
294        linker.site_value(
295            compute_cuda_site_symbol(),
296            DefaultFactory.opaque(Arc::new(site))?,
297        )?;
298        Ok(())
299    }
300}
301
302fn cuda_input(
303    runtime: &Arc<CudaLibrarySet>,
304    tensor: &Tensor,
305) -> std::result::Result<Arc<CudaDeviceBuffer>, TensorExecError> {
306    if let Some(buffer) = tensor
307        .storage()
308        .as_any()
309        .downcast_ref::<CudaResidentStorage>()
310        .and_then(CudaResidentStorage::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("cuda 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: CudaLoadError) -> TensorExecError {
328    TensorExecError::Eval {
329        message: Arc::from(error.to_string()),
330    }
331}
sim_lib_compute_cuda/site.rs

sim_lib_compute_cuda/
site.rs
