rlx_wgpu/
device.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4//! wgpu device discovery + capabilities.
5//!
6//! [`wgpu_device`] returns a process-global singleton. [`select_vulkan_backend`]
7//! routes subsequent calls to a Vulkan-only instance (for [`Device::Vulkan`]).
8
9use std::sync::OnceLock;
10use std::sync::atomic::{AtomicU8, Ordering};
11
12const PREF_DEFAULT: u8 = 0;
13const PREF_VULKAN: u8 = 1;
14
15static BACKEND_PREF: AtomicU8 = AtomicU8::new(PREF_DEFAULT);
16
17/// Detected wgpu adapter + device. We hold them together because
18/// every command submission needs both the device (for encoding) and
19/// the queue (for committing).
20pub struct WgpuDevice {
21    pub instance: wgpu::Instance,
22    pub adapter: wgpu::Adapter,
23    pub device: wgpu::Device,
24    pub queue: wgpu::Queue,
25    pub name: String,
26    pub backend: wgpu::Backend,
27}
28
29impl WgpuDevice {
30    fn new_with_backends(backends: wgpu::Backends) -> Option<Self> {
31        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
32            backends,
33            flags: wgpu::InstanceFlags::default(),
34            backend_options: wgpu::BackendOptions::default(),
35            memory_budget_thresholds: wgpu::MemoryBudgetThresholds::default(),
36            display: None,
37        });
38        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
39            power_preference: wgpu::PowerPreference::HighPerformance,
40            compatible_surface: None,
41            force_fallback_adapter: false,
42        }))
43        .ok()?;
44
45        let info = adapter.get_info();
46        let limits = adapter.limits();
47        let adapter_feats = adapter.features();
48        let mut required_features = wgpu::Features::empty();
49        if adapter_feats.contains(wgpu::Features::SHADER_F16) {
50            required_features |= wgpu::Features::SHADER_F16;
51        }
52        if adapter_feats.contains(wgpu::Features::EXPERIMENTAL_COOPERATIVE_MATRIX) {
53            required_features |= wgpu::Features::EXPERIMENTAL_COOPERATIVE_MATRIX;
54        }
55        if adapter_feats.contains(wgpu::Features::SUBGROUP) {
56            required_features |= wgpu::Features::SUBGROUP;
57        }
58
59        let (device, queue) =
60            match pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
61                label: Some("rlx-wgpu device"),
62                required_features,
63                required_limits: limits,
64                memory_hints: wgpu::MemoryHints::Performance,
65                experimental_features: unsafe { wgpu::ExperimentalFeatures::enabled() },
66                trace: wgpu::Trace::Off,
67            })) {
68                Ok(p) => p,
69                Err(e) => {
70                    eprintln!("rlx-wgpu request_device failed: {e}");
71                    return None;
72                }
73            };
74
75        Some(Self {
76            instance,
77            adapter,
78            device,
79            queue,
80            name: info.name,
81            backend: info.backend,
82        })
83    }
84
85    fn new_default() -> Option<Self> {
86        Self::new_with_backends(wgpu::Backends::from_env().unwrap_or(wgpu::Backends::all()))
87    }
88}
89
90// SAFETY: wgpu's Device + Queue are documented thread-safe.
91unsafe impl Send for WgpuDevice {}
92unsafe impl Sync for WgpuDevice {}
93
94fn default_device() -> Option<&'static WgpuDevice> {
95    static DEVICE: OnceLock<Option<WgpuDevice>> = OnceLock::new();
96    DEVICE.get_or_init(WgpuDevice::new_default).as_ref()
97}
98
99fn vulkan_device() -> Option<&'static WgpuDevice> {
100    static DEVICE: OnceLock<Option<WgpuDevice>> = OnceLock::new();
101    DEVICE
102        .get_or_init(|| WgpuDevice::new_with_backends(wgpu::Backends::VULKAN))
103        .as_ref()
104}
105
106/// Prefer the Vulkan-only wgpu instance for [`Device::Vulkan`] sessions.
107/// Call before the first [`wgpu_device`] use in that process (or use
108/// `Device::Vulkan` via the runtime registry, which calls this).
109pub fn select_vulkan_backend() {
110    BACKEND_PREF.store(PREF_VULKAN, Ordering::SeqCst);
111}
112
113/// True when a Vulkan adapter is reachable (MoltenVK on macOS, native on Linux/Windows).
114pub fn is_vulkan_available() -> bool {
115    vulkan_device().is_some()
116}
117
118/// Get or initialize the global wgpu device singleton. Returns None
119/// on systems with no compatible adapter.
120pub fn wgpu_device() -> Option<&'static WgpuDevice> {
121    if BACKEND_PREF.load(Ordering::SeqCst) == PREF_VULKAN {
122        vulkan_device()
123    } else {
124        default_device()
125    }
126}
rlx_wgpu/device.rs

rlx_wgpu/
device.rs
