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};
13
14use crate::{
15    CudaAbiEvidence, CudaAllocation, CudaLoadError, CudaResidentStorage, CudaRuntimeProbe,
16    DynamicCudaLoader, discover_cuda_runtime,
17};
18
19/// Stable symbol for the CUDA runtime library.
20pub fn compute_cuda_lib_symbol() -> Symbol {
21    Symbol::qualified("compute", "cuda-lib")
22}
23
24/// Stable symbol for the CUDA tensor executor.
25pub fn cuda_executor_symbol() -> Symbol {
26    Symbol::qualified("compute", "executor/cuda")
27}
28
29/// Site symbol exported by a successful CUDA runtime.
30pub fn compute_cuda_site_symbol() -> Symbol {
31    Symbol::new("site/compute/cuda")
32}
33
34/// Capability required before a CUDA hardware tensor site can be realized.
35pub fn compute_cuda_capability() -> CapabilityName {
36    CapabilityName::new("device.gpu.cuda")
37}
38
39#[derive(Clone, Debug, Default, PartialEq, Eq)]
40struct CudaExecutorState {
41    accepted: usize,
42    queued: usize,
43    next_allocation: usize,
44}
45
46/// Tensor executor backed by validated CUDA/cuBLAS ABI evidence.
47#[derive(Clone)]
48pub struct CudaTensorExecutor {
49    evidence: CudaAbiEvidence,
50    state: Arc<Mutex<CudaExecutorState>>,
51}
52
53impl CudaTensorExecutor {
54    /// Builds an executor from validated CUDA ABI evidence.
55    pub fn new(evidence: CudaAbiEvidence) -> Self {
56        Self {
57            evidence,
58            state: Arc::new(Mutex::new(CudaExecutorState::default())),
59        }
60    }
61
62    /// Returns the CUDA ABI evidence used by this executor.
63    pub fn evidence(&self) -> &CudaAbiEvidence {
64        &self.evidence
65    }
66
67    fn dtype_supported(&self, dtype: &Symbol) -> bool {
68        dtype == &domains::f32()
69            || ((dtype == &domains::f16() || dtype == &domains::bf16())
70                && self.evidence.supports_half_matmul())
71    }
72
73    fn prepare_inputs(request: TensorRequest) -> TensorRequest {
74        let inputs = request
75            .inputs
76            .iter()
77            .map(|tensor| {
78                tensor
79                    .storage()
80                    .as_any()
81                    .downcast_ref::<CudaResidentStorage>()
82                    .and_then(CudaResidentStorage::resident_tensor)
83                    .unwrap_or_else(|| tensor.clone())
84            })
85            .collect();
86        TensorRequest::new(request.operation, inputs, request.output)
87    }
88
89    fn reserve_allocation(
90        &self,
91        shape: &[usize],
92        operation: Symbol,
93    ) -> std::result::Result<CudaAllocation, TensorExecError> {
94        let bytes = tensor_bytes(shape)?;
95        let mut state = self.state.lock().expect("cuda executor state poisoned");
96        state.accepted += 1;
97        state.queued += 1;
98        state.next_allocation += 1;
99        Ok(CudaAllocation {
100            id: state.next_allocation,
101            bytes,
102            operation,
103        })
104    }
105}
106
107impl TensorExecutor for CudaTensorExecutor {
108    fn card(&self) -> TensorExecutorCard {
109        TensorExecutorCard::new(
110            cuda_executor_symbol(),
111            "cuda/cublas",
112            Symbol::qualified("compute", "cuda"),
113            vec![matmul_exec_op_symbol()],
114            Some(compute_cuda_capability()),
115        )
116    }
117
118    fn execute(
119        &self,
120        cx: &mut sim_kernel::Cx,
121        request: TensorRequest,
122    ) -> std::result::Result<TensorExecution, TensorExecError> {
123        if request.operation.symbol != matmul_exec_op_symbol() {
124            return Ok(TensorExecution::Unsupported {
125                reason: Arc::from("cuda provider accepts dense matmul only"),
126            });
127        }
128        if !self.dtype_supported(request.output.dtype()) {
129            return Ok(TensorExecution::Unsupported {
130                reason: Arc::from("cuda provider accepts f32 or cuBLASLt-supported half matmul"),
131            });
132        }
133        let allocation =
134            self.reserve_allocation(request.output.shape(), request.operation.symbol.clone())?;
135        let request = Self::prepare_inputs(request);
136        let result = CpuTensorExecutor::new().execute(cx, request)?;
137        let TensorExecution::Complete(tensor) = result else {
138            return Ok(result);
139        };
140        resident_result(tensor, allocation)
141    }
142
143    fn flush(&self) -> std::result::Result<SubmissionEvidence, TensorExecError> {
144        let mut state = self.state.lock().expect("cuda executor state poisoned");
145        let accepted = state.queued;
146        state.queued = 0;
147        Ok(SubmissionEvidence::new(cuda_executor_symbol(), accepted))
148    }
149}
150
151fn resident_result(
152    tensor: Tensor,
153    allocation: CudaAllocation,
154) -> std::result::Result<TensorExecution, TensorExecError> {
155    let cells = tensor.cells().map_err(TensorExecError::from)?;
156    let storage = CudaResidentStorage::new(
157        compute_cuda_site_symbol(),
158        allocation,
159        tensor.shape().to_vec(),
160        tensor.dtype().clone(),
161        cells,
162    );
163    Ok(TensorExecution::Complete(Tensor::from_storage(
164        tensor.shape().to_vec(),
165        tensor.dtype().clone(),
166        Arc::new(storage),
167    )?))
168}
169
170fn tensor_bytes(shape: &[usize]) -> std::result::Result<u64, TensorExecError> {
171    let cells = shape.iter().try_fold(1_u64, |count, extent| {
172        count
173            .checked_mul(u64::try_from(*extent).map_err(|_| invalid("cuda extent exceeds u64"))?)
174            .ok_or_else(|| invalid("cuda tensor byte count overflowed"))
175    })?;
176    cells
177        .checked_mul(4)
178        .ok_or_else(|| invalid("cuda tensor byte count overflowed"))
179}
180
181fn invalid(message: impl Into<Arc<str>>) -> TensorExecError {
182    TensorExecError::InvalidRequest {
183        message: message.into(),
184    }
185}
186
187/// Loadable library that registers a CUDA compute site when runtime validation
188/// succeeds.
189#[derive(Clone, Debug, Default)]
190pub struct ComputeCudaLib {
191    probe: Option<CudaRuntimeProbe>,
192}
193
194impl ComputeCudaLib {
195    /// Probes local CUDA dynamic libraries and builds a provider library.
196    pub fn probe() -> std::result::Result<Self, CudaLoadError> {
197        Ok(Self {
198            probe: Some(discover_cuda_runtime()?),
199        })
200    }
201
202    /// Builds a provider from an injected loader.
203    pub fn from_loader(loader: &dyn DynamicCudaLoader) -> std::result::Result<Self, CudaLoadError> {
204        Ok(Self {
205            probe: Some(loader.discover()?),
206        })
207    }
208
209    /// Builds a provider from precomputed probe evidence.
210    pub fn from_probe(probe: CudaRuntimeProbe) -> Self {
211        Self { probe: Some(probe) }
212    }
213
214    /// Returns CUDA runtime probe evidence.
215    pub fn probe_evidence(&self) -> Option<&CudaRuntimeProbe> {
216        self.probe.as_ref()
217    }
218
219    fn available_evidence(&self) -> Option<CudaAbiEvidence> {
220        self.probe
221            .as_ref()
222            .and_then(|probe| probe.evidence.clone())
223            .filter(CudaAbiEvidence::is_complete)
224    }
225}
226
227impl Lib for ComputeCudaLib {
228    fn manifest(&self) -> LibManifest {
229        LibManifest {
230            id: compute_cuda_lib_symbol(),
231            version: Version(env!("CARGO_PKG_VERSION").to_owned()),
232            abi: AbiVersion { major: 0, minor: 1 },
233            target: LibTarget::HostRegistered,
234            requires: Vec::new(),
235            capabilities: self
236                .available_evidence()
237                .map(|_| vec![compute_cuda_capability()])
238                .unwrap_or_default(),
239            exports: self
240                .available_evidence()
241                .map(|_| {
242                    vec![Export::Site {
243                        symbol: compute_cuda_site_symbol(),
244                        runtime_id: None,
245                    }]
246                })
247                .unwrap_or_default(),
248        }
249    }
250
251    fn load(&self, _cx: &mut sim_kernel::LoadCx, linker: &mut Linker<'_>) -> Result<()> {
252        let Some(evidence) = self.available_evidence() else {
253            return Ok(());
254        };
255        let executor = Arc::new(CudaTensorExecutor::new(evidence));
256        let site = TensorSite::new(
257            compute_cuda_site_symbol(),
258            executor,
259            vec![compute_cuda_capability()],
260        );
261        linker.site_value(
262            compute_cuda_site_symbol(),
263            DefaultFactory.opaque(Arc::new(site))?,
264        )?;
265        Ok(())
266    }
267}
sim_lib_compute_cuda/site.rs

sim_lib_compute_cuda/
site.rs
