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

sim_lib_compute_wgpu/
site.rs
