1use crate::library::PipelineLibrary;
2use crate::pipeline::PipelineHandle;
3use crate::{ComputeProgram, GpuRuntime, WarmupBudget, WarmupProgress};
4use std::fmt::{self, Display, Formatter};
5use std::sync::atomic::{AtomicUsize, Ordering};
6use std::sync::{Arc, Mutex};
7use std::time::Instant;
8use wgpu::{
9 Adapter, AdapterInfo, Backend, Backends, Device, DeviceLostReason, DeviceType, Features,
10 Limits, PowerPreference, Queue,
11};
12
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14pub enum LimitsPolicy {
15 Minimum,
16
17 Adapter,
18}
19
20pub struct GpuRequest {
21 pub backends: Backends,
22 pub power_preference: PowerPreference,
23
24 pub device_name: Option<String>,
25
26 pub required_features: Features,
27
28 pub optional_features: Features,
29 pub limits: LimitsPolicy,
30}
31
32impl Default for GpuRequest {
33 fn default() -> Self {
34 Self {
35 backends: GpuRequest::NATIVE_BACKENDS,
36 power_preference: PowerPreference::HighPerformance,
37 device_name: None,
38 required_features: Features::empty(),
39 optional_features: GpuRequest::PROFILING_FEATURES,
40 limits: LimitsPolicy::Adapter,
41 }
42 }
43}
44
45impl GpuRequest {
46 #[cfg(feature = "profile")]
47 pub const PROFILING_FEATURES: Features =
48 Features::TIMESTAMP_QUERY.union(Features::TIMESTAMP_QUERY_INSIDE_ENCODERS);
49 #[cfg(not(feature = "profile"))]
50 pub const PROFILING_FEATURES: Features = Features::empty();
51
52 pub const NATIVE_BACKENDS: Backends = Backends::DX12
53 .union(Backends::METAL)
54 .union(Backends::VULKAN);
55
56 pub fn adapter_named(name: impl Into<String>) -> Self {
57 Self {
58 device_name: Some(name.into()),
59 ..Self::default()
60 }
61 }
62
63 pub fn minimum_limits() -> Self {
64 Self {
65 limits: LimitsPolicy::Minimum,
66 ..Self::default()
67 }
68 }
69}
70
71#[derive(Clone, Debug, PartialEq, Eq)]
72pub enum GpuUnavailable {
73 NoAdapter {
74 backends: Backends,
75 },
76 DeviceNotFound {
77 requested: String,
78 available: Vec<String>,
79 },
80 MissingFeatures {
81 missing: Features,
82 available: Features,
83 },
84 DeviceRejected {
85 message: String,
86 },
87 InsufficientLimits {
88 adapter: String,
89 },
90}
91
92impl Display for GpuUnavailable {
93 fn fmt(&self, out: &mut Formatter<'_>) -> fmt::Result {
94 match self {
95 Self::NoAdapter { backends } => {
96 write!(
97 out,
98 "no adapter exposed by the enabled backends {backends:?}"
99 )
100 }
101 Self::DeviceNotFound {
102 requested,
103 available,
104 } => write!(
105 out,
106 "no adapter name contains {requested:?}; available: {available:?}"
107 ),
108 Self::MissingFeatures { missing, available } => write!(
109 out,
110 "adapter lacks required features {missing:?} (offers {available:?})"
111 ),
112 Self::DeviceRejected { message } => write!(out, "device request rejected: {message}"),
113 Self::InsufficientLimits { adapter } => write!(
114 out,
115 "adapter {adapter:?} cannot satisfy the storage-buffer limits dynamis requires"
116 ),
117 }
118 }
119}
120
121impl std::error::Error for GpuUnavailable {}
122
123#[derive(Clone, Debug, PartialEq, Eq)]
124pub struct DeviceLost {
125 pub reason: DeviceLostReason,
126 pub message: String,
127}
128
129impl Display for DeviceLost {
130 fn fmt(&self, out: &mut Formatter<'_>) -> fmt::Result {
131 write!(out, "{:?}: {}", self.reason, self.message)
132 }
133}
134
135#[derive(Default)]
136struct Health {
137 lost: Mutex<Option<DeviceLost>>,
138}
139
140pub struct GpuContext {
141 adapter: Adapter,
142 device: Device,
143 queue: Queue,
144 features: Features,
145 limits: Limits,
146 timestamp_period_ns: f32,
147 health: Arc<Health>,
148 pipelines: Arc<Mutex<PipelineLibrary>>,
149 compilation: Arc<Mutex<()>>,
150}
151
152impl GpuContext {
153 pub const MINIMUM_LIMITS: Limits = Limits {
154 max_storage_buffers_per_shader_stage: 16,
155 max_storage_textures_per_shader_stage: 0,
156 max_storage_buffer_binding_size: 16 * 1024 * 1024,
157 max_buffer_size: 64 * 1024 * 1024,
158 max_compute_workgroup_size_x: 256,
159 max_compute_workgroup_size_y: 64,
160 max_compute_workgroup_size_z: 64,
161 max_compute_workgroups_per_dimension: 65_535,
162 ..Limits::defaults()
163 };
164
165 pub async fn open(request: &GpuRequest) -> Result<Self, GpuUnavailable> {
166 let adapter = select_adapter(request).await?;
167 let available = adapter.features();
168 let missing = request.required_features.difference(available);
169 if !missing.is_empty() {
170 return Err(GpuUnavailable::MissingFeatures { missing, available });
171 }
172 let features = request.required_features | (request.optional_features & available);
173 let limits = match request.limits {
174 LimitsPolicy::Adapter => adapter.limits(),
175 LimitsPolicy::Minimum => Self::MINIMUM_LIMITS,
176 };
177 if !Self::MINIMUM_LIMITS.check_limits(&adapter.limits()) {
178 return Err(GpuUnavailable::InsufficientLimits {
179 adapter: adapter.get_info().name,
180 });
181 }
182 let (device, queue) = adapter
183 .request_device(&wgpu::DeviceDescriptor {
184 label: Some("dynamis device"),
185 required_features: features,
186 required_limits: limits.clone(),
187 memory_hints: wgpu::MemoryHints::Performance,
188 ..Default::default()
189 })
190 .await
191 .map_err(|error| GpuUnavailable::DeviceRejected {
192 message: error.to_string(),
193 })?;
194 let health = Arc::new(Health::default());
195 let callback = health.clone();
196 device.set_device_lost_callback(move |reason, message| {
197 *callback.lost.lock().unwrap() = Some(DeviceLost { reason, message });
198 });
199 let timestamp_period_ns = queue.get_timestamp_period();
200 Ok(Self {
201 adapter,
202 device,
203 queue,
204 features,
205 limits,
206 timestamp_period_ns,
207 health,
208 pipelines: Arc::new(Mutex::new(PipelineLibrary::new())),
209 compilation: Arc::new(Mutex::new(())),
210 })
211 }
212
213 pub async fn new() -> Self {
214 match Self::open(&GpuRequest::default()).await {
215 Ok(context) => context,
216 Err(error) => panic!("{error}"),
217 }
218 }
219
220 pub async fn available_adapters(backends: Backends) -> Vec<AdapterInfo> {
221 GpuRuntime::shared()
222 .await
223 .adapters(backends)
224 .iter()
225 .map(|adapter| adapter.get_info())
226 .collect()
227 }
228
229 pub fn device(&self) -> &Device {
230 &self.device
231 }
232
233 pub fn queue(&self) -> &Queue {
234 &self.queue
235 }
236
237 pub fn adapter_info(&self) -> AdapterInfo {
238 self.adapter.get_info()
239 }
240
241 pub fn features(&self) -> Features {
242 self.features
243 }
244
245 pub fn workgroups_per_row(&self) -> u32 {
246 self.limits.max_compute_workgroups_per_dimension.max(1)
247 }
248
249 pub fn limits(&self) -> &Limits {
250 &self.limits
251 }
252
253 pub fn supports(&self, features: Features) -> bool {
254 self.features.contains(features)
255 }
256
257 pub fn timestamp_period_ns(&self) -> f32 {
258 self.timestamp_period_ns
259 }
260
261 #[cfg(feature = "profile")]
262 pub fn supports_pass_timing(&self) -> bool {
263 self.supports(GpuRequest::PROFILING_FEATURES)
264 }
265
266 pub fn device_lost(&self) -> Option<DeviceLost> {
267 self.health.lost.lock().unwrap().clone()
268 }
269
270 pub fn assert_alive(&self) {
271 if let Some(lost) = self.device_lost() {
272 panic!("gpu device lost: {lost}");
273 }
274 }
275
276 pub fn poll(&self) {
277 self.assert_alive();
278 self.device
279 .poll(wgpu::PollType::Poll)
280 .expect("GPU polling failed");
281 self.assert_alive();
282 }
283
284 pub fn declare(&self, program: ComputeProgram) -> PipelineHandle {
285 self.pipelines
286 .lock()
287 .unwrap()
288 .declare(&self.device, program)
289 }
290
291 pub fn warmup(&self, budget: WarmupBudget) -> WarmupProgress {
292 let _compiling = self.compilation.lock().unwrap();
293 match budget {
294 WarmupBudget::All => self.compile_all_pending(),
295 WarmupBudget::Within(duration) => {
296 let deadline = Instant::now().checked_add(duration);
297 loop {
298 let Some(pending) = self.pipelines.lock().unwrap().take_pending() else {
299 break;
300 };
301 pending.compile(&self.device);
302 if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
303 break;
304 }
305 }
306 }
307 }
308 self.pipelines.lock().unwrap().progress()
309 }
310
311 fn compile_all_pending(&self) {
312 let pending = self.pipelines.lock().unwrap().drain_pending();
313 if pending.is_empty() {
314 return;
315 }
316 let workers = std::thread::available_parallelism()
317 .map(|threads| threads.get())
318 .unwrap_or(1)
319 .min(pending.len());
320 let next = AtomicUsize::new(0);
321 std::thread::scope(|scope| {
322 for _ in 0..workers {
323 scope.spawn(|| {
324 loop {
325 let index = next.fetch_add(1, Ordering::Relaxed);
326 let Some(pipeline) = pending.get(index) else {
327 break;
328 };
329 pipeline.compile(&self.device);
330 }
331 });
332 }
333 });
334 }
335
336 pub fn is_warm(&self) -> bool {
337 self.pipelines.lock().unwrap().progress().complete()
338 }
339}
340
341impl Clone for GpuContext {
342 fn clone(&self) -> Self {
343 Self {
344 adapter: self.adapter.clone(),
345 device: self.device.clone(),
346 queue: self.queue.clone(),
347 features: self.features,
348 limits: self.limits.clone(),
349 timestamp_period_ns: self.timestamp_period_ns,
350 health: self.health.clone(),
351 pipelines: self.pipelines.clone(),
352 compilation: self.compilation.clone(),
353 }
354 }
355}
356
357async fn select_adapter(request: &GpuRequest) -> Result<Adapter, GpuUnavailable> {
358 let adapters = GpuRuntime::shared().await.adapters(request.backends);
359 if adapters.is_empty() {
360 return Err(GpuUnavailable::NoAdapter {
361 backends: request.backends,
362 });
363 }
364 let candidates = match &request.device_name {
365 Some(needle) => {
366 let wanted = needle.to_lowercase();
367 let matched = adapters
368 .iter()
369 .filter(|adapter| adapter.get_info().name.to_lowercase().contains(&wanted))
370 .cloned()
371 .collect::<Vec<_>>();
372 if matched.is_empty() {
373 return Err(GpuUnavailable::DeviceNotFound {
374 requested: needle.clone(),
375 available: adapters
376 .iter()
377 .map(|adapter| adapter.get_info().name)
378 .collect(),
379 });
380 }
381 matched
382 }
383 None => adapters,
384 };
385 Ok(prefer(candidates, request.power_preference))
386}
387
388fn prefer(candidates: Vec<Adapter>, power: PowerPreference) -> Adapter {
389 let mut candidates = candidates;
390 let rank = |adapter: &Adapter| {
391 let info = adapter.get_info();
392 let device = match power {
393 PowerPreference::HighPerformance => match info.device_type {
394 DeviceType::DiscreteGpu => 0u8,
395 DeviceType::VirtualGpu => 1,
396 DeviceType::IntegratedGpu => 2,
397 _ => 3,
398 },
399 PowerPreference::LowPower => match info.device_type {
400 DeviceType::IntegratedGpu => 0u8,
401 DeviceType::VirtualGpu => 1,
402 DeviceType::DiscreteGpu => 2,
403 _ => 3,
404 },
405 PowerPreference::None => 0u8,
406 };
407 (native_backend_rank(info.backend), device)
408 };
409 candidates.sort_by_key(|adapter| rank(adapter));
410 candidates
411 .into_iter()
412 .next()
413 .unwrap_or_else(|| panic!("adapter candidates must be non-empty"))
414}
415
416fn native_backend_rank(backend: Backend) -> u8 {
417 #[cfg(target_os = "windows")]
418 const PRIMARY: Backend = Backend::Dx12;
419 #[cfg(any(target_os = "macos", target_os = "ios", target_os = "visionos"))]
420 const PRIMARY: Backend = Backend::Metal;
421 #[cfg(not(any(
422 target_os = "windows",
423 target_os = "macos",
424 target_os = "ios",
425 target_os = "visionos"
426 )))]
427 const PRIMARY: Backend = Backend::Vulkan;
428
429 u8::from(backend != PRIMARY)
430}