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