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