1use std::future::Future;
2use std::pin::Pin;
3use std::sync::{Arc, OnceLock};
4use std::task::{Context, Poll, Wake, Waker};
5use std::thread::{self, Thread};
6use vyre_driver::error::{Error, Result};
7
8use crate::staging_reserve::reserve_backend_vec;
9
10#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
16pub struct EnabledFeatures {
17 pub timestamp_query: bool,
19 pub timestamp_query_inside_encoders: bool,
21 pub subgroup: bool,
23 pub subgroup_barrier: bool,
25 pub shader_f16: bool,
27 pub pipeline_cache: bool,
29 pub push_constants: bool,
31 pub indirect_first_instance: bool,
33 pub max_workgroup_size: [u32; 3],
35 pub max_storage_buffer_binding_size: u64,
37 pub max_subgroup_size: u32,
39 pub min_subgroup_size: u32,
44}
45
46pub(crate) fn poll_device_once(
47 device: &wgpu::Device,
48) -> std::result::Result<wgpu::PollStatus, vyre_driver::BackendError> {
49 device.poll(wgpu::PollType::Poll).map_err(|error| {
50 vyre_driver::BackendError::new(format!(
51 "wgpu device poll failed: {error}. Fix: inspect device loss and driver health before reusing this backend."
52 ))
53 })
54}
55
56pub(crate) fn poll_device_wait_for(
57 device: &wgpu::Device,
58 submission: wgpu::SubmissionIndex,
59) -> std::result::Result<wgpu::PollStatus, vyre_driver::BackendError> {
60 device
61 .poll(wgpu::PollType::wait_for(submission))
62 .map_err(|error| {
63 vyre_driver::BackendError::new(format!(
64 "wgpu device wait-for-submission poll failed: {error}. Fix: inspect device loss, driver health, and submission lifetime before reusing this backend."
65 ))
66 })
67}
68
69struct CachedRuntime {
70 device_queue: Arc<(wgpu::Device, wgpu::Queue)>,
71 adapter_info: wgpu::AdapterInfo,
72 #[cfg(test)]
73 enabled_features: EnabledFeatures,
74}
75
76static CACHED_RUNTIME: OnceLock<Result<CachedRuntime>> = OnceLock::new();
77
78fn cached_runtime() -> &'static Result<CachedRuntime> {
79 CACHED_RUNTIME.get_or_init(|| {
80 #[cfg(test)]
81 let ((device, queue), adapter_info, enabled_features) = init_device()?;
82 #[cfg(not(test))]
83 let ((device, queue), adapter_info, _enabled_features) = init_device()?;
84 Ok(CachedRuntime {
85 device_queue: Arc::new((device, queue)),
86 adapter_info,
87 #[cfg(test)]
88 enabled_features,
89 })
90 })
91}
92
93#[inline]
118pub fn cached_device() -> Result<Arc<(wgpu::Device, wgpu::Queue)>> {
119 cached_runtime()
120 .as_ref()
121 .map(|runtime| Arc::clone(&runtime.device_queue))
122 .map_err(Clone::clone)
123}
124
125#[inline]
131pub fn cached_adapter_info() -> Result<&'static wgpu::AdapterInfo> {
132 cached_runtime()
133 .as_ref()
134 .map(|runtime| &runtime.adapter_info)
135 .map_err(Clone::clone)
136}
137
138#[cfg(test)]
140pub(crate) fn cached_enabled_features() -> Result<&'static EnabledFeatures> {
141 cached_runtime()
142 .as_ref()
143 .map(|runtime| &runtime.enabled_features)
144 .map_err(Clone::clone)
145}
146
147#[cfg(test)]
149#[inline]
150pub(crate) fn is_cached_device(device: &wgpu::Device) -> bool {
151 CACHED_RUNTIME
152 .get()
153 .and_then(|res| res.as_ref().ok())
154 .map(|runtime| &runtime.device_queue.0 == device)
155 .unwrap_or(false)
156}
157
158#[inline]
165pub fn init_device() -> Result<(
166 (wgpu::Device, wgpu::Queue),
167 wgpu::AdapterInfo,
168 EnabledFeatures,
169)> {
170 let gpu = wait_for_gpu(acquire_gpu())?;
171 Ok(gpu)
172}
173
174#[inline]
181pub async fn acquire_gpu() -> Result<(
182 (wgpu::Device, wgpu::Queue),
183 wgpu::AdapterInfo,
184 EnabledFeatures,
185)> {
186 if let Some(index) = super::selector::adapter_index_from_env()? {
187 return super::selector::acquire_gpu_for_adapter(index).await;
188 }
189
190 let instance = wgpu::Instance::default();
191 let adapters = instance.enumerate_adapters(wgpu::Backends::all());
192 let mut candidates = Vec::new();
193 reserve_probe_vec(
194 &mut candidates,
195 adapters.len(),
196 "GPU acquisition candidates",
197 )?;
198 candidates.extend(adapters.iter().filter_map(|adapter| {
199 let info = adapter.get_info();
200 crate::capabilities::is_real_gpu(&info).then(|| {
201 let score = gpu_candidate_score(&info, adapter.features(), &adapter.limits());
202 (adapter, info, score)
203 })
204 }));
205 candidates.sort_by(|left, right| right.2.cmp(&left.2));
206
207 let mut failures = Vec::new();
208 reserve_probe_vec(&mut failures, candidates.len(), "GPU acquisition failures")?;
209 for (adapter, info, _) in candidates {
210 match request_device_for_adapter(adapter, "vyre device").await {
211 Ok(device) => return Ok(device),
212 Err(error) => failures.push(format!("{} ({:?}): {error}", info.name, info.device_type)),
213 }
214 }
215
216 let mut probed = Vec::new();
217 reserve_probe_vec(&mut probed, adapters.len(), "GPU acquisition probe report")?;
218 probed.extend(adapters.iter().map(|adapter| {
219 let info = adapter.get_info();
220 format!(
221 "{} ({:?}, backend={:?})",
222 info.name, info.device_type, info.backend
223 )
224 }));
225 Err(Error::Gpu {
226 message: format!(
227 "no real GPU adapter could create a wgpu device. Probed adapters: [{}]. Device failures: [{}]. Fix: expose a discrete, integrated, or virtual GPU through a wgpu-supported driver before running vyre.",
228 probed.join(", "),
229 failures.join("; ")
230 ),
231 })
232}
233
234pub(super) async fn request_device_for_adapter(
235 adapter: &wgpu::Adapter,
236 label: &'static str,
237) -> Result<(
238 (wgpu::Device, wgpu::Queue),
239 wgpu::AdapterInfo,
240 EnabledFeatures,
241)> {
242 let adapter_info = adapter.get_info();
243 if !crate::capabilities::is_real_gpu(&adapter_info) {
244 return Err(Error::Gpu {
245 message: format!(
246 "wgpu adapter `{}` reports device type {:?}, which is not a real GPU execution target. Fix: select a discrete, integrated, or virtual GPU adapter; CPU/software adapters are not production dispatch backends.",
247 adapter_info.name, adapter_info.device_type
248 ),
249 });
250 }
251 let adapter_features = adapter.features();
261 let adapter_limits = adapter.limits();
262 let (features, mut enabled) =
263 enabled_features_for_adapter(adapter_features, &adapter_limits, adapter_info.backend);
264
265 let device_queue = adapter
266 .request_device(
267 &wgpu::DeviceDescriptor {
268 label: Some(label),
269 required_features: features,
270 required_limits: wgpu::Limits {
271 max_compute_workgroup_size_x: adapter_limits.max_compute_workgroup_size_x,
272 max_compute_workgroup_size_y: adapter_limits.max_compute_workgroup_size_y,
273 max_compute_workgroup_size_z: adapter_limits.max_compute_workgroup_size_z,
274 max_compute_invocations_per_workgroup: adapter_limits
275 .max_compute_invocations_per_workgroup,
276 max_compute_workgroups_per_dimension: adapter_limits
277 .max_compute_workgroups_per_dimension,
278 max_compute_workgroup_storage_size: adapter_limits
279 .max_compute_workgroup_storage_size,
280 max_storage_buffer_binding_size: adapter_limits.max_storage_buffer_binding_size,
281 max_buffer_size: adapter_limits.max_buffer_size,
289 min_subgroup_size: if enabled.subgroup {
290 adapter_limits.min_subgroup_size
291 } else {
292 0
293 },
294 max_subgroup_size: if enabled.subgroup {
295 adapter_limits.max_subgroup_size
296 } else {
297 0
298 },
299 max_storage_buffers_per_shader_stage:
300 adapter_limits.max_storage_buffers_per_shader_stage,
301 max_push_constant_size: if enabled.push_constants {
302 adapter_limits.max_push_constant_size
303 } else {
304 0
305 },
306 ..wgpu::Limits::default()
307 },
308 memory_hints: wgpu::MemoryHints::default(),
309 trace: wgpu::Trace::Off,
310 },
311 )
312 .await
313 .map_err(|error| Error::Gpu {
314 message: format!("failed to acquire device for adapter `{}`: {error}. Fix: check requested wgpu limits/features against the adapter and update the GPU driver if limits are unexpectedly low.", adapter_info.name),
315 })?;
316 let device_limits = device_queue.0.limits();
317 enabled.max_workgroup_size = [
318 device_limits.max_compute_workgroup_size_x,
319 device_limits.max_compute_workgroup_size_y,
320 device_limits.max_compute_workgroup_size_z,
321 ];
322 enabled.max_storage_buffer_binding_size =
323 u64::from(device_limits.max_storage_buffer_binding_size);
324 enabled.max_subgroup_size = device_limits.max_subgroup_size;
325 enabled.min_subgroup_size = device_limits.min_subgroup_size;
326
327 if enabled.subgroup {
328 subgroup_smoke_compiles(&device_queue.0).map_err(|error| Error::Gpu {
329 message: format!(
330 "adapter `{}` advertises SUBGROUP but rejects the subgroup compute-pipeline smoke test: {error}. Fix: repair the wgpu feature negotiation or GPU driver; do not silently report subgroup support as disabled on a subgroup-capable adapter.",
331 adapter_info.name
332 ),
333 })?;
334 }
335
336 Ok((device_queue, adapter_info, enabled))
337}
338
339fn backend_implements_pipeline_cache(backend: wgpu::Backend) -> bool {
346 matches!(backend, wgpu::Backend::Vulkan | wgpu::Backend::Dx12)
347}
348
349pub(super) fn enabled_features_for_adapter(
350 adapter_features: wgpu::Features,
351 adapter_limits: &wgpu::Limits,
352 backend: wgpu::Backend,
353) -> (wgpu::Features, EnabledFeatures) {
354 let mut features = wgpu::Features::empty();
355 let mut enabled = EnabledFeatures::default();
356 if adapter_features.contains(wgpu::Features::TIMESTAMP_QUERY) {
357 features |= wgpu::Features::TIMESTAMP_QUERY;
358 enabled.timestamp_query = true;
359 }
360 if adapter_features.contains(wgpu::Features::TIMESTAMP_QUERY_INSIDE_ENCODERS) {
361 features |= wgpu::Features::TIMESTAMP_QUERY_INSIDE_ENCODERS;
362 enabled.timestamp_query = true;
363 enabled.timestamp_query_inside_encoders = true;
364 }
365 if crate::capabilities::supports_subgroup_for_adapter(adapter_features, adapter_limits) {
366 features |= wgpu::Features::SUBGROUP;
367 enabled.subgroup = true;
368 }
369 if adapter_features.contains(wgpu::Features::SUBGROUP_BARRIER) {
370 features |= wgpu::Features::SUBGROUP_BARRIER;
371 enabled.subgroup_barrier = true;
372 }
373 if adapter_features.contains(wgpu::Features::SHADER_F16) {
374 features |= wgpu::Features::SHADER_F16;
375 enabled.shader_f16 = true;
376 }
377 if adapter_features.contains(wgpu::Features::PIPELINE_CACHE)
378 && backend_implements_pipeline_cache(backend)
379 {
380 features |= wgpu::Features::PIPELINE_CACHE;
381 enabled.pipeline_cache = true;
382 }
383 if adapter_features.contains(wgpu::Features::PUSH_CONSTANTS) {
384 features |= wgpu::Features::PUSH_CONSTANTS;
385 enabled.push_constants = true;
386 }
387 if adapter_features.contains(wgpu::Features::INDIRECT_FIRST_INSTANCE) {
388 features |= wgpu::Features::INDIRECT_FIRST_INSTANCE;
389 enabled.indirect_first_instance = true;
390 }
391
392 enabled.max_workgroup_size = [
393 adapter_limits.max_compute_workgroup_size_x,
394 adapter_limits.max_compute_workgroup_size_y,
395 adapter_limits.max_compute_workgroup_size_z,
396 ];
397 enabled.max_storage_buffer_binding_size =
398 u64::from(adapter_limits.max_storage_buffer_binding_size);
399 enabled.max_subgroup_size = adapter_limits.max_subgroup_size;
400 enabled.min_subgroup_size = adapter_limits.min_subgroup_size;
401 (features, enabled)
402}
403
404fn real_gpu_rank(device_type: wgpu::DeviceType) -> u8 {
405 match device_type {
406 wgpu::DeviceType::DiscreteGpu => 3,
407 wgpu::DeviceType::IntegratedGpu => 2,
408 wgpu::DeviceType::VirtualGpu => 1,
409 wgpu::DeviceType::Cpu | wgpu::DeviceType::Other => 0,
410 }
411}
412
413fn gpu_candidate_score(
414 info: &wgpu::AdapterInfo,
415 adapter_features: wgpu::Features,
416 adapter_limits: &wgpu::Limits,
417) -> u128 {
418 let mut feature_score = 0u128;
419 if crate::capabilities::supports_subgroup_for_adapter(adapter_features, adapter_limits) {
420 feature_score |= 1 << 7;
421 }
422 if adapter_features.contains(wgpu::Features::SUBGROUP_BARRIER) {
423 feature_score |= 1 << 6;
424 }
425 if adapter_features.contains(wgpu::Features::SHADER_F16) {
426 feature_score |= 1 << 5;
427 }
428 if adapter_features.contains(wgpu::Features::PIPELINE_CACHE) {
429 feature_score |= 1 << 4;
430 }
431 if adapter_features.contains(wgpu::Features::PUSH_CONSTANTS) {
432 feature_score |= 1 << 3;
433 }
434 if adapter_features.contains(wgpu::Features::INDIRECT_FIRST_INSTANCE) {
435 feature_score |= 1 << 2;
436 }
437 if adapter_features.contains(wgpu::Features::TIMESTAMP_QUERY) {
438 feature_score |= 1 << 1;
439 }
440 if adapter_features.contains(wgpu::Features::TIMESTAMP_QUERY_INSIDE_ENCODERS) {
441 feature_score |= 1;
442 }
443
444 let storage_binding_bits = u128::from(adapter_limits.max_storage_buffer_binding_size.ilog2());
445 let buffer_bits = u128::from(adapter_limits.max_buffer_size.max(1).ilog2());
446 let workgroup_invocations = u128::from(adapter_limits.max_compute_invocations_per_workgroup);
447 let workgroup_storage_bits = u128::from(
448 adapter_limits
449 .max_compute_workgroup_storage_size
450 .max(1)
451 .ilog2(),
452 );
453 let storage_buffers = u128::from(adapter_limits.max_storage_buffers_per_shader_stage);
454
455 (u128::from(real_gpu_rank(info.device_type)) << 120)
456 | (feature_score << 96)
457 | (storage_binding_bits << 88)
458 | (buffer_bits << 80)
459 | (workgroup_invocations << 56)
460 | (workgroup_storage_bits << 48)
461 | storage_buffers
462}
463
464fn subgroup_smoke_compiles(device: &wgpu::Device) -> std::result::Result<(), String> {
465 const WGSL: &str = r#"
466@compute @workgroup_size(32)
467fn main(@builtin(subgroup_invocation_id) lane: u32, @builtin(subgroup_size) size: u32) {
468 let total = subgroupAdd(lane + size);
469 if (total == 0u) {
470 return;
471 }
472}
473"#;
474
475 device.push_error_scope(wgpu::ErrorFilter::Validation);
476 let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
477 label: Some("vyre subgroup capability probe"),
478 source: wgpu::ShaderSource::Wgsl(WGSL.into()),
479 });
480 let _pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
481 label: Some("vyre subgroup capability probe"),
482 layout: None,
483 module: &module,
484 entry_point: Some("main"),
485 compilation_options: wgpu::PipelineCompilationOptions::default(),
486 cache: None,
487 });
488 match pop_error_scope_now(device) {
489 Ok(None) => Ok(()),
490 Ok(Some(error)) => Err(format!("validation error: {error}")),
491 Err(error) => Err(error.to_string()),
492 }
493}
494
495struct ThreadWaker(Thread);
496
497impl Wake for ThreadWaker {
498 fn wake(self: Arc<Self>) {
499 self.0.unpark();
500 }
501
502 fn wake_by_ref(self: &Arc<Self>) {
503 self.0.unpark();
504 }
505}
506
507struct NoopWaker;
508
509impl Wake for NoopWaker {
510 fn wake(self: Arc<Self>) {}
511
512 fn wake_by_ref(self: &Arc<Self>) {}
513}
514
515pub(crate) fn pop_error_scope_now(
516 device: &wgpu::Device,
517) -> std::result::Result<Option<wgpu::Error>, &'static str> {
518 device
519 .poll(wgpu::PollType::Poll)
520 .map_err(|_| "wgpu device poll failed before error-scope pop")?;
521 let waker = Waker::from(Arc::new(NoopWaker));
522 let mut context = Context::from_waker(&waker);
523 let mut future = Box::pin(device.pop_error_scope());
524 match Future::poll(Pin::as_mut(&mut future), &mut context) {
525 Poll::Ready(error) => Ok(error),
526 Poll::Pending => Err(
527 "wgpu error scope did not resolve after a nonblocking device poll. Fix: inspect the backend event loop; validation must not require a hot-path host wait.",
528 ),
529 }
530}
531
532pub(super) fn wait_for_gpu<T>(future: impl Future<Output = T>) -> T {
533 let waker = Waker::from(Arc::new(ThreadWaker(thread::current())));
534 let mut context = Context::from_waker(&waker);
535 let mut future = Box::pin(future);
536 loop {
537 match Pin::as_mut(&mut future).poll(&mut context) {
538 Poll::Ready(value) => return value,
539 Poll::Pending => thread::park(),
540 }
541 }
542}
543
544fn reserve_probe_vec<T>(vec: &mut Vec<T>, additional: usize, context: &'static str) -> Result<()> {
545 reserve_backend_vec(vec, additional, context).map_err(|error| Error::Gpu {
546 message: error.to_string(),
547 })
548}
549
550#[cfg(test)]
551mod tests {
552 use super::*;
553
554 #[test]
556 fn cached_device_is_singleton() {
557 let first = cached_device().expect("Fix: GPU must be available for runtime tests");
558 let second = cached_device().expect("Fix: GPU must be available for runtime tests");
559 assert!(
560 Arc::ptr_eq(&first, &second),
561 "cached_device must return the same Arc after singleton initialization"
562 );
563 assert!(
564 is_cached_device(&first.0),
565 "legacy shared APIs must still recognize cached_device-created devices"
566 );
567 }
568
569 #[test]
570 fn cached_adapter_info_uses_cached_runtime() {
571 let info = cached_adapter_info().expect("Fix: cached adapter info must share GPU init");
572 let enabled =
573 cached_enabled_features().expect("Fix: cached runtime must retain capability snapshot");
574 let device_queue = cached_device().expect("Fix: GPU must be available for runtime tests");
575 assert!(
576 !info.name.is_empty(),
577 "cached adapter info must come from the initialized runtime adapter"
578 );
579 assert!(
580 enabled.max_workgroup_size.iter().all(|axis| *axis > 0),
581 "cached runtime must retain nonzero device workgroup limits for capability reporting"
582 );
583 assert!(
584 is_cached_device(&device_queue.0),
585 "cached adapter info must not replace the cached device with a second init"
586 );
587 }
588
589 #[test]
590 fn gpu_candidate_score_prefers_stronger_compute_adapter_within_same_class() {
591 let info = wgpu::AdapterInfo {
592 name: "gpu".to_string(),
593 vendor: 0x10de,
594 device: 0x2c02,
595 device_type: wgpu::DeviceType::DiscreteGpu,
596 driver: "nvidia".to_string(),
597 driver_info: "test".to_string(),
598 backend: wgpu::Backend::Vulkan,
599 };
600 let weak_limits = wgpu::Limits {
601 max_storage_buffer_binding_size: 1 << 20,
602 max_buffer_size: 1 << 28,
603 max_compute_invocations_per_workgroup: 256,
604 max_compute_workgroup_storage_size: 16 << 10,
605 max_storage_buffers_per_shader_stage: 8,
606 ..wgpu::Limits::default()
607 };
608 let strong_limits = wgpu::Limits {
609 max_storage_buffer_binding_size: 1 << 30,
610 max_buffer_size: 1 << 34,
611 max_compute_invocations_per_workgroup: 1024,
612 max_compute_workgroup_storage_size: 64 << 10,
613 max_storage_buffers_per_shader_stage: 16,
614 min_subgroup_size: 32,
615 max_subgroup_size: 32,
616 ..wgpu::Limits::default()
617 };
618 let weak = gpu_candidate_score(&info, wgpu::Features::empty(), &weak_limits);
619 let strong = gpu_candidate_score(
620 &info,
621 wgpu::Features::SUBGROUP
622 | wgpu::Features::SUBGROUP_BARRIER
623 | wgpu::Features::SHADER_F16
624 | wgpu::Features::PIPELINE_CACHE,
625 &strong_limits,
626 );
627
628 assert!(
629 strong > weak,
630 "Fix: automatic GPU acquisition must prefer the stronger same-class compute adapter."
631 );
632 }
633
634 #[test]
635 fn production_device_acquisition_uses_fallible_probe_reservations() {
636 let production = include_str!("device.rs")
637 .split("\n#[cfg(test)]\nmod tests")
638 .next()
639 .expect("Fix: device production section should precede tests");
640
641 assert!(
642 !production.contains("Vec::with_capacity"),
643 "Fix: centralized GPU acquisition must not use infallible capacity constructors."
644 );
645 assert!(
646 production.contains("reserve_probe_vec"),
647 "Fix: centralized GPU acquisition should reserve probe metadata through the shared staging helper."
648 );
649 assert!(
650 production.contains("reserve_backend_vec"),
651 "Fix: WGPU device acquisition should reuse the backend staging reservation policy."
652 );
653 }
654}