sim_lib_compute_wgpu/
site.rs

1//! Loadable wgpu 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    SubmissionEvidence, TensorExecError, TensorExecution, TensorExecutor, TensorExecutorCard,
11    TensorRequest, TensorSite, domains,
12};
13
14use crate::{
15    ProbePolicy, WgpuAdapterProbe, WgpuDiscovery, WgpuKernelDType, WgpuPhysicalCounters,
16    WgpuPipelineCache, WgpuProbePort, WgpuQueueLimits, WgpuResidentArena, WgpuResidentStorage,
17    WgpuResidentStorageDescriptor, WgpuSegmentPlan, WgpuTileProfile,
18    dispatch::{execute_pointwise_dispatch, is_pointwise_dispatch},
19    dispatch_linalg::execute_linalg_dispatch,
20    dispatch_reductions::execute_reduction_dispatch,
21    kernels::{execute_portable_kernel, kernel_op},
22};
23
24/// Stable symbol for the wgpu runtime library.
25pub fn compute_wgpu_lib_symbol() -> Symbol {
26    Symbol::qualified("compute", "wgpu-lib")
27}
28
29/// Stable symbol for a wgpu tensor executor.
30pub fn wgpu_executor_symbol(ordinal: usize) -> Symbol {
31    Symbol::qualified("compute", format!("executor/wgpu/{ordinal}"))
32}
33
34/// Site symbol exported for a successful wgpu adapter.
35pub fn compute_wgpu_site_symbol(ordinal: usize) -> Symbol {
36    Symbol::new(format!("site/compute/wgpu/{ordinal}"))
37}
38
39/// Capability required before a wgpu hardware tensor site can be realized.
40pub fn compute_wgpu_capability() -> CapabilityName {
41    CapabilityName::new("device.gpu.wgpu")
42}
43
44/// Tensor executor descriptor backed by a successful wgpu probe.
45#[derive(Clone)]
46pub struct WgpuTensorExecutor {
47    pub(crate) probe: WgpuAdapterProbe,
48    pub(crate) state: Arc<Mutex<WgpuExecutorState>>,
49    pub(crate) context: Option<WgpuExecutionContext>,
50}
51
52#[derive(Debug)]
53pub(crate) struct WgpuExecutorState {
54    pub(crate) pipelines: WgpuPipelineCache,
55    arena: WgpuResidentArena,
56    queued: usize,
57    queued_bytes: u64,
58    accepted: usize,
59    physical: WgpuPhysicalCounters,
60}
61
62#[derive(Clone, Debug)]
63pub(crate) struct WgpuExecutionContext {
64    pub(crate) device: Arc<wgpu::Device>,
65    pub(crate) queue: Arc<wgpu::Queue>,
66}
67
68impl WgpuTensorExecutor {
69    /// Builds an executor from successful probe evidence.
70    pub fn new(probe: WgpuAdapterProbe) -> Self {
71        Self::from_parts(probe, None)
72    }
73
74    pub(crate) fn from_parts(
75        probe: WgpuAdapterProbe,
76        context: Option<WgpuExecutionContext>,
77    ) -> Self {
78        let arena_bytes = probe.adapter.granted_limits.max_buffer_size.max(4);
79        Self {
80            probe,
81            state: Arc::new(Mutex::new(WgpuExecutorState {
82                pipelines: WgpuPipelineCache::default(),
83                arena: WgpuResidentArena::new(arena_bytes),
84                queued: 0,
85                queued_bytes: 0,
86                accepted: 0,
87                physical: WgpuPhysicalCounters::default(),
88            })),
89            context,
90        }
91    }
92
93    /// Returns the underlying probe evidence.
94    pub fn probe(&self) -> &WgpuAdapterProbe {
95        &self.probe
96    }
97
98    /// Returns the pipeline cache snapshot.
99    pub fn pipeline_cache_snapshot(&self) -> crate::WgpuPipelineCacheSnapshot {
100        self.state
101            .lock()
102            .expect("wgpu executor state poisoned")
103            .pipelines
104            .snapshot()
105    }
106
107    /// Returns queue-derived physical submission evidence.
108    pub fn physical_evidence(&self) -> crate::PhysicalSubmissionEvidence {
109        self.state
110            .lock()
111            .expect("wgpu executor state poisoned")
112            .physical
113            .snapshot()
114    }
115
116    pub(crate) fn physical_counters(&self) -> WgpuPhysicalCounters {
117        self.state
118            .lock()
119            .expect("wgpu executor state poisoned")
120            .physical
121            .clone()
122    }
123
124    fn dtype_for(
125        &self,
126        request: &TensorRequest,
127    ) -> std::result::Result<WgpuKernelDType, TensorExecError> {
128        let dtype = request.output.dtype();
129        if dtype == &domains::f32() || dtype == &domains::f64() {
130            Ok(WgpuKernelDType::F32)
131        } else if dtype == &domains::f16() {
132            if self.probe.adapter.granted_features.shader_f16 {
133                Ok(WgpuKernelDType::F16Native)
134            } else {
135                Ok(WgpuKernelDType::Bf16WidenedToF32)
136            }
137        } else if dtype == &domains::bf16() {
138            Ok(WgpuKernelDType::Bf16WidenedToF32)
139        } else {
140            Err(unsupported(
141                request.operation.symbol.clone(),
142                "wgpu portable kernels accept f32/f64/half-family tensor dtypes",
143            ))
144        }
145    }
146
147    fn check_submission_limits(&self, bytes: u64) -> std::result::Result<(), TensorExecError> {
148        let state = self.state.lock().expect("wgpu executor state poisoned");
149        let tile = WgpuTileProfile::from_probe(&self.probe);
150        let limits = WgpuQueueLimits {
151            max_nodes: 64,
152            max_bytes: tile.max_dispatch_bytes,
153            deadline_tick: u64::MAX,
154        };
155        if state.queued >= limits.max_nodes {
156            return Err(invalid("wgpu submission queue node limit reached"));
157        }
158        if state.queued_bytes.saturating_add(bytes) > limits.max_bytes {
159            return Err(invalid("wgpu submission queue byte limit reached"));
160        }
161        Ok(())
162    }
163}
164
165impl TensorExecutor for WgpuTensorExecutor {
166    fn card(&self) -> TensorExecutorCard {
167        TensorExecutorCard::new(
168            wgpu_executor_symbol(self.probe.adapter.ordinal),
169            format!(
170                "wgpu/{}/{}",
171                self.probe.adapter.backend, self.probe.adapter.name
172            ),
173            Symbol::qualified("compute", "wgpu"),
174            vec![
175                sim_lib_numbers_tensor::add_op_symbol(),
176                sim_lib_numbers_tensor::sub_op_symbol(),
177                sim_lib_numbers_tensor::mul_op_symbol(),
178                sim_lib_numbers_tensor::div_op_symbol(),
179                sim_lib_numbers_tensor::neg_op_symbol(),
180                sim_lib_numbers_tensor::sqrt_op_symbol(),
181                sim_lib_numbers_tensor::exp_op_symbol(),
182                Symbol::qualified("tensor", "op/log"),
183                sim_lib_numbers_tensor::sin_op_symbol(),
184                sim_lib_numbers_tensor::cos_op_symbol(),
185                sim_lib_numbers_tensor::sum_op_symbol(),
186                sim_lib_numbers_tensor::min_op_symbol(),
187                sim_lib_numbers_tensor::max_op_symbol(),
188                sim_lib_numbers_tensor::norm_op_symbol(),
189                sim_lib_numbers_tensor::transpose_exec_op_symbol(),
190                sim_lib_numbers_tensor::dot_op_symbol(),
191                sim_lib_numbers_tensor::matmul_exec_op_symbol(),
192            ],
193            Some(compute_wgpu_capability()),
194        )
195    }
196
197    fn execute(
198        &self,
199        cx: &mut sim_kernel::Cx,
200        request: TensorRequest,
201    ) -> std::result::Result<TensorExecution, TensorExecError> {
202        let Some(op) = kernel_op(&request.operation.symbol) else {
203            return Ok(TensorExecution::Unsupported {
204                reason: Arc::from("operation is outside the portable wgpu kernel set"),
205            });
206        };
207        let dtype = self.dtype_for(&request)?;
208        let bytes = tensor_bytes(request.output.shape())?;
209        self.check_submission_limits(bytes)?;
210        let dispatched = if is_pointwise_dispatch(op) {
211            Some(execute_pointwise_dispatch(
212                self, cx, &request, op, dtype, bytes,
213            )?)
214        } else if op.is_reduction() {
215            Some(execute_reduction_dispatch(self, cx, &request, op, dtype)?)
216        } else if op.is_linalg() {
217            Some(execute_linalg_dispatch(self, cx, &request, op, dtype)?)
218        } else {
219            None
220        };
221        let (buffer, pipeline_symbol, len) = if let Some(output) = dispatched {
222            (output.buffer, Some(output.pipeline), output.len)
223        } else {
224            let tensor = execute_portable_kernel(cx, &request, dtype)?;
225            let values = crate::dispatch::tensor_f32_values(cx, &tensor, dtype)?;
226            let bytes = crate::dispatch::f32_bytes(&values);
227            let Some(context) = &self.context else {
228                return Err(invalid("wgpu device context is unavailable"));
229            };
230            let buffer = context.device.create_buffer(&wgpu::BufferDescriptor {
231                label: Some("sim-compute-wgpu-portable-upload"),
232                size: bytes.len().max(4) as u64,
233                usage: wgpu::BufferUsages::STORAGE
234                    | wgpu::BufferUsages::COPY_DST
235                    | wgpu::BufferUsages::COPY_SRC,
236                mapped_at_creation: false,
237            });
238            context.queue.write_buffer(&buffer, 0, &bytes);
239            self.physical_counters().record_upload(bytes.len() as u64);
240            (Arc::new(buffer), None, values.len())
241        };
242        let boundary = self
243            .probe
244            .adapter
245            .granted_limits
246            .max_storage_buffer_binding_size
247            .max(4);
248        let segments = WgpuSegmentPlan::new(bytes, boundary, boundary);
249        self.physical_counters().record_submit(&segments.segments);
250        let pipeline = {
251            let mut state = self.state.lock().expect("wgpu executor state poisoned");
252            let allocation = state.arena.allocate(bytes.max(4)).map_err(invalid)?;
253            state.queued += 1;
254            state.queued_bytes += bytes;
255            state.accepted += 1;
256            let pipeline = if let Some(pipeline_symbol) = pipeline_symbol {
257                pipeline_symbol
258            } else {
259                state
260                    .pipelines
261                    .get_or_insert(&self.probe, op, dtype, request.output.shape().len())
262                    .symbol
263            };
264            (allocation, pipeline)
265        };
266        let storage = WgpuResidentStorage::new(WgpuResidentStorageDescriptor {
267            site: compute_wgpu_site_symbol(self.probe.adapter.ordinal),
268            allocation: pipeline.0,
269            pipeline: pipeline.1,
270            segments: segments.segments,
271            dtype: request.output.dtype().clone(),
272            len,
273            buffer,
274            context: self
275                .context
276                .clone()
277                .ok_or_else(|| invalid("wgpu device context is unavailable"))?,
278            counters: self.physical_counters(),
279        });
280        Ok(TensorExecution::Complete(
281            sim_lib_numbers_tensor::Tensor::from_storage(
282                request.output.shape().to_vec(),
283                request.output.dtype().clone(),
284                Arc::new(storage),
285            )?,
286        ))
287    }
288
289    fn flush(&self) -> std::result::Result<SubmissionEvidence, TensorExecError> {
290        let mut state = self.state.lock().expect("wgpu executor state poisoned");
291        let accepted = state.queued;
292        state.queued = 0;
293        state.queued_bytes = 0;
294        Ok(SubmissionEvidence::new(
295            wgpu_executor_symbol(self.probe.adapter.ordinal),
296            accepted,
297        ))
298    }
299}
300
301fn tensor_bytes(shape: &[usize]) -> std::result::Result<u64, TensorExecError> {
302    let cells = shape.iter().try_fold(1_u64, |count, extent| {
303        count
304            .checked_mul(
305                u64::try_from(*extent).map_err(|_| invalid("wgpu tensor extent exceeds u64"))?,
306            )
307            .ok_or_else(|| invalid("wgpu tensor byte count overflowed"))
308    })?;
309    cells
310        .checked_mul(4)
311        .ok_or_else(|| invalid("wgpu tensor byte count overflowed"))
312}
313
314fn invalid(message: impl Into<Arc<str>>) -> TensorExecError {
315    TensorExecError::InvalidRequest {
316        message: message.into(),
317    }
318}
319
320fn unsupported(operation: Symbol, reason: impl Into<Arc<str>>) -> TensorExecError {
321    TensorExecError::Unsupported {
322        operation,
323        reason: reason.into(),
324    }
325}
326
327/// Loadable library that registers only successful wgpu adapter sites.
328#[derive(Clone, Debug, Default)]
329pub struct ComputeWgpuLib {
330    discovery: WgpuDiscovery,
331    contexts: Vec<WgpuExecutionContext>,
332}
333
334impl ComputeWgpuLib {
335    /// Builds a library from a capsule-owned probe port.
336    pub fn from_probe_port(port: &dyn WgpuProbePort, policy: &ProbePolicy) -> Result<Self> {
337        let runtimes = port
338            .probe_wgpu(policy)
339            .map_err(|err| sim_kernel::Error::Eval(err.to_string()))?;
340        let mut probes = Vec::with_capacity(runtimes.len());
341        let mut contexts = Vec::with_capacity(runtimes.len());
342        for runtime in runtimes {
343            probes.push(runtime.probe);
344            contexts.push(WgpuExecutionContext {
345                device: Arc::new(runtime.device),
346                queue: Arc::new(runtime.queue),
347            });
348        }
349        let discovery = WgpuDiscovery::from_probes(probes, Vec::new());
350        Ok(Self {
351            discovery,
352            contexts,
353        })
354    }
355
356    /// Builds a library from precomputed discovery evidence.
357    pub fn from_discovery(discovery: WgpuDiscovery) -> Self {
358        Self {
359            discovery,
360            contexts: Vec::new(),
361        }
362    }
363
364    /// Returns discovery evidence, including failed-adapter diagnostics.
365    pub fn discovery(&self) -> &WgpuDiscovery {
366        &self.discovery
367    }
368}
369
370impl Lib for ComputeWgpuLib {
371    fn manifest(&self) -> LibManifest {
372        LibManifest {
373            id: compute_wgpu_lib_symbol(),
374            version: Version(env!("CARGO_PKG_VERSION").to_owned()),
375            abi: AbiVersion { major: 0, minor: 1 },
376            target: LibTarget::HostRegistered,
377            requires: Vec::new(),
378            capabilities: if self.discovery.adapters.is_empty() {
379                Vec::new()
380            } else {
381                vec![compute_wgpu_capability()]
382            },
383            exports: self
384                .discovery
385                .adapters
386                .iter()
387                .map(|probe| Export::Site {
388                    symbol: compute_wgpu_site_symbol(probe.adapter.ordinal),
389                    runtime_id: None,
390                })
391                .collect(),
392        }
393    }
394
395    fn load(&self, _cx: &mut sim_kernel::LoadCx, linker: &mut Linker<'_>) -> Result<()> {
396        for probe in &self.discovery.adapters {
397            let symbol = compute_wgpu_site_symbol(probe.adapter.ordinal);
398            let executor = if let Some(context) = self.contexts.get(probe.adapter.ordinal) {
399                Arc::new(WgpuTensorExecutor::from_parts(
400                    probe.clone(),
401                    Some(context.clone()),
402                ))
403            } else {
404                Arc::new(WgpuTensorExecutor::new(probe.clone()))
405            };
406            let site = TensorSite::new(symbol.clone(), executor, vec![compute_wgpu_capability()]);
407            linker.site_value(symbol, DefaultFactory.opaque(Arc::new(site))?)?;
408        }
409        Ok(())
410    }
411}
sim_lib_compute_wgpu/site.rs

sim_lib_compute_wgpu/
site.rs
