Skip to main content

cubecl_wgpu/
runtime.rs

1use std::marker::PhantomData;
2
3use crate::WgpuCompiler;
4use crate::{
5    AutoCompiler, AutoGraphicsApi, GraphicsApi, WgpuBackend, WgpuDevice, WgpuDeviceKind, backend,
6    compute::WgpuServer, contiguous_strides,
7};
8use cubecl_common::device::{Device, DeviceService, ServiceId};
9use cubecl_common::profile::TimingMethod;
10use cubecl_core::WgpuCompilationOptions;
11use cubecl_core::device::{DeviceId, ServerUtilitiesHandle};
12use cubecl_core::ir::TargetProperties;
13use cubecl_core::server::ServerUtilities;
14use cubecl_core::zspace::{Shape, Strides};
15use cubecl_environment::{future, sync::Mutex};
16use cubecl_ir::{DeviceIdentity, DeviceProperties, HardwareProperties, MemoryDeviceProperties};
17use cubecl_monitoring::{
18    DevicePower, DeviceTemperature, DeviceUtilization, UtilizationUnavailable,
19    gpu_utilization::CardCounters,
20};
21use cubecl_server::allocator::ContiguousMemoryLayoutPolicy;
22#[cfg(not(feature = "vulkan-validate"))]
23use cubecl_server::logging::ProfileLevel;
24pub use cubecl_server::memory_management::MemoryConfiguration;
25use cubecl_server::runtime::{DeviceDescription, DeviceKind, Runtime};
26use cubecl_server::{client::Client, logging::ServerLogger};
27use wgpu::{InstanceFlags, RequestAdapterOptions};
28
29/// Runtime that uses the [wgpu] crate.
30///
31/// The default [`AutoCompiler`] selects a native shader compiler when supported and falls back to
32/// WGSL otherwise. Supply an explicit compiler type when fallback is not desired.
33/// With `msl` enabled, [`AutoCompiler`] requires native MSL support for devices pinned to Metal;
34/// an automatic device can still fall back to WGSL. An explicit [`crate::WgslCompiler`] continues
35/// to use WGSL on Metal without requiring native MSL support.
36///
37/// For advanced configuration, use [`init_setup`] to pass runtime options or select a specific
38/// graphics API.
39#[derive(Debug)]
40pub struct WgpuRuntime<Compiler = AutoCompiler> {
41    _p: PhantomData<Compiler>,
42}
43
44impl<C> Clone for WgpuRuntime<C> {
45    fn clone(&self) -> Self {
46        Self { _p: self._p }
47    }
48}
49
50impl<C: WgpuCompiler> DeviceService for WgpuServer<C> {
51    fn init(device_id: cubecl_common::device::DeviceId) -> Self {
52        let device = WgpuDevice::from_id(device_id);
53        let setup =
54            future::block_on(try_create_setup(&device)).expect("Unable to acquire wgpu device");
55        try_create_server(setup, RuntimeOptions::default(), device_id, device.backend)
56            .expect("Unable to initialize wgpu runtime")
57    }
58
59    fn utilities(&self) -> ServerUtilitiesHandle {
60        self.utilities.clone() as ServerUtilitiesHandle
61    }
62}
63
64impl<C: WgpuCompiler> Runtime for WgpuRuntime<C> {
65    type Server = WgpuServer<C>;
66    type Device = WgpuDevice;
67
68    fn can_read_tensor(shape: &Shape, strides: &Strides) -> bool {
69        if shape.is_empty() {
70            return true;
71        }
72
73        for (&expected, &stride) in contiguous_strides(shape).iter().zip(strides.iter()) {
74            if expected != stride {
75                return false;
76            }
77        }
78
79        true
80    }
81
82    /// The properties with no device to pick a compiler by. A server takes its own compiler's
83    /// ([`WgpuCompiler::target_properties`]), which differ under MSL.
84    fn target_properties() -> TargetProperties {
85        TargetProperties::default()
86    }
87
88    fn enumerate_devices(type_id: u16) -> Vec<DeviceId> {
89        // The devices of that type on `Auto`, whose id is index zero with no
90        // graphics API in the top bits.
91        Self::enumerate_devices_like(DeviceId::new(type_id, 0))
92    }
93
94    fn is_available() -> bool {
95        // A software rasterizer — lavapipe, llvmpipe, WARP — enumerates as
96        // `WgpuDeviceKind::Cpu`. It runs, but a machine with nothing else is
97        // better served by a native CPU runtime, so wgpu does not claim it;
98        // a caller who wants it still names it.
99        let gpu = [
100            WgpuDeviceKind::DiscreteGpu(0),
101            WgpuDeviceKind::IntegratedGpu(0),
102            WgpuDeviceKind::VirtualGpu(0),
103            WgpuDeviceKind::Other(0),
104        ]
105        .map(|kind| WgpuDevice::new(kind).to_id().type_id);
106
107        Self::enumerate_all_devices()
108            .iter()
109            .any(|device| gpu.contains(&device.type_id))
110    }
111
112    /// Each adapter once, on whichever graphics API [`WgpuBackend::Auto`]
113    /// settles on — which is what a client created lazily resolves these ids
114    /// against. The same adapter pinned to another API is a device a caller
115    /// names, not one more to count.
116    fn enumerate_all_devices() -> Vec<DeviceId> {
117        adapters_on(WgpuBackend::Auto)
118    }
119
120    fn enumerate_devices_like(device_id: DeviceId) -> Vec<DeviceId> {
121        let device = WgpuDevice::from_id(device_id);
122        let reachable = adapters_on(device.backend);
123
124        match device.kind {
125            // One per graphics API, standing for whichever adapter it lands
126            // on: its own only peer, there as soon as the API has anything.
127            // Listing the adapters beside it would hand one of them a second
128            // client under another id.
129            WgpuDeviceKind::DefaultDevice if !reachable.is_empty() => alloc::vec![device_id],
130            _ => reachable
131                .into_iter()
132                .filter(|id| id.type_id == device_id.type_id)
133                .collect(),
134        }
135    }
136
137    /// The browser hands out one adapter without saying what it is, so a kind
138    /// there is only a power preference — and `request_adapter` honors the
139    /// low-power one as well as the rest.
140    #[cfg(target_family = "wasm")]
141    fn find_device(device_id: DeviceId) -> Result<(), usize> {
142        let device = WgpuDevice::from_id(device_id);
143        let peers = Self::enumerate_devices_like(device_id);
144
145        let low_power = device.kind == WgpuDeviceKind::IntegratedGpu(0)
146            && !adapters_on(device.backend).is_empty();
147
148        match peers.contains(&device_id) || low_power {
149            true => Ok(()),
150            false => Err(peers.len()),
151        }
152    }
153
154    fn utilization(device: &Self::Device) -> Result<DeviceUtilization, UtilizationUnavailable> {
155        CardCounters::utilization(Self::client(device).properties().identity.physical.as_ref())
156    }
157
158    fn temperature(device: &Self::Device) -> Result<DeviceTemperature, UtilizationUnavailable> {
159        CardCounters::temperature(Self::client(device).properties().identity.physical.as_ref())
160    }
161
162    fn power(device: &Self::Device) -> Result<DevicePower, UtilizationUnavailable> {
163        CardCounters::power(Self::client(device).properties().identity.physical.as_ref())
164    }
165
166    /// The adapter description, from the list its graphics API enumerates: the one `request_adapter` would
167    /// select, in the order `enumerate_devices` lists them. It reports no memory size.
168    fn description(device: &Self::Device) -> Option<DeviceDescription> {
169        #[cfg(target_family = "wasm")]
170        {
171            let _ = device;
172            None
173        }
174        #[cfg(not(target_family = "wasm"))]
175        {
176            let (kind, index) = match device.kind {
177                WgpuDeviceKind::DiscreteGpu(index) => (wgpu::DeviceType::DiscreteGpu, index),
178                WgpuDeviceKind::IntegratedGpu(index) => (wgpu::DeviceType::IntegratedGpu, index),
179                WgpuDeviceKind::VirtualGpu(index) => (wgpu::DeviceType::VirtualGpu, index),
180                WgpuDeviceKind::Other(index) => (wgpu::DeviceType::Other, index),
181                WgpuDeviceKind::Cpu => (wgpu::DeviceType::Cpu, 0),
182                // The adapter it lands on is chosen at creation.
183                WgpuDeviceKind::DefaultDevice | WgpuDeviceKind::Existing(_) => return None,
184            };
185            let (api, _) = settle(device.backend)?;
186            probe_with(api, |adapters| {
187                let adapter = adapters
188                    .iter()
189                    .filter(|adapter| adapter.get_info().device_type == kind)
190                    .nth(index)?;
191                Some(describe_adapter(adapter))
192            })
193        }
194    }
195}
196
197/// The ids of the adapters `backend` reaches, pinned the way it is.
198///
199/// Those of the one graphics API it settles on, so each id resolves in
200/// `WgpuServer::init` to the adapter it was listed for. A device brought up on
201/// another API through [`init_setup`] is reached through the client that call
202/// hands back, not here.
203fn adapters_on(backend: WgpuBackend) -> Vec<DeviceId> {
204    // WebGPU only supports a single device currently, and only the browser's
205    // own API reaches it.
206    #[cfg(target_family = "wasm")]
207    let ids = match backend {
208        WgpuBackend::Auto | WgpuBackend::WebGpu => vec![DeviceId::new(0, 0)],
209        _ => Vec::new(),
210    };
211
212    #[cfg(not(target_family = "wasm"))]
213    let ids = settle(backend)
214        .map(|(_, adapters)| adapter_device_ids(adapters))
215        .unwrap_or_default();
216
217    ids.into_iter()
218        .map(|id| WgpuDevice::from_id(id).on(backend).to_id())
219        .collect()
220}
221
222/// The graphics API `backend` settles on, and the adapters this machine has
223/// there: the first of its candidates with a GPU, or failing that, the first
224/// with anything at all.
225///
226/// A software rasterizer is not reason enough to stop. Where Vulkan offers
227/// only lavapipe and `OpenGL` the real GPU — a VM passing it through, say —
228/// stopping at Vulkan leaves that GPU unreachable through `Auto`, and wgpu
229/// declining a machine it could have served.
230#[cfg(not(target_family = "wasm"))]
231fn settle(backend: WgpuBackend) -> Option<(wgpu::Backend, Vec<wgpu::AdapterInfo>)> {
232    let mut software_only = None;
233
234    for api in backend_candidates(backend) {
235        let adapters = probe(api);
236
237        if adapters
238            .iter()
239            .any(|adapter| adapter.device_type != wgpu::DeviceType::Cpu)
240        {
241            return Some((api, adapters));
242        }
243
244        if software_only.is_none() && !adapters.is_empty() {
245            software_only = Some((api, adapters));
246        }
247    }
248
249    software_only
250}
251
252/// The `wgpu` backends to try for a [`WgpuBackend`], best first.
253///
254/// A pinned one is the only candidate — that is what pinning it means.
255fn backend_candidates(backend: WgpuBackend) -> alloc::vec::Vec<wgpu::Backend> {
256    match backend {
257        WgpuBackend::Auto => AutoGraphicsApi::chain(),
258        WgpuBackend::Vulkan => alloc::vec![wgpu::Backend::Vulkan],
259        WgpuBackend::Metal => alloc::vec![wgpu::Backend::Metal],
260        WgpuBackend::Dx12 => alloc::vec![wgpu::Backend::Dx12],
261        WgpuBackend::Gl => alloc::vec![wgpu::Backend::Gl],
262        WgpuBackend::WebGpu => alloc::vec![wgpu::Backend::BrowserWebGpu],
263    }
264}
265
266/// The instance each backend's devices come up on, kept so a device going away never destroys
267/// one. Locked whenever an instance is created or dropped: the Vulkan loader crashes when
268/// instances come and go on several threads at once.
269static INSTANCES: Mutex<Vec<(wgpu::Backend, wgpu::Instance)>> = Mutex::new(Vec::new());
270
271/// The adapters `backend` has on this machine, asked through an instance limited to it.
272#[cfg(not(target_family = "wasm"))]
273fn probe(backend: wgpu::Backend) -> Vec<wgpu::AdapterInfo> {
274    probe_with(backend, |adapters| {
275        adapters.iter().map(wgpu::Adapter::get_info).collect()
276    })
277}
278
279/// What `adapter` says about its hardware.
280#[cfg(not(target_family = "wasm"))]
281fn describe_adapter(adapter: &wgpu::Adapter) -> DeviceDescription {
282    let info = adapter.get_info();
283    let kind = match info.device_type {
284        wgpu::DeviceType::DiscreteGpu => DeviceKind::Discrete,
285        wgpu::DeviceType::IntegratedGpu => DeviceKind::Integrated,
286        wgpu::DeviceType::VirtualGpu => DeviceKind::Virtual,
287        wgpu::DeviceType::Cpu => DeviceKind::Cpu,
288        wgpu::DeviceType::Other => DeviceKind::Other,
289    };
290    let mut description = DeviceDescription::new(info.name.clone(), kind);
291    description.physical = backend::physical_device(adapter, &info);
292    // The driver's name and its own version string, e.g. `NVIDIA` and `580.82.09`; either can
293    // be empty.
294    let driver = [info.driver.trim(), info.driver_info.trim()]
295        .into_iter()
296        .filter(|part| !part.is_empty())
297        .collect::<Vec<_>>()
298        .join(" ");
299    description.driver_version = (!driver.is_empty()).then_some(driver);
300    description
301}
302
303/// What `inspect` makes of the adapters `backend` has on this machine.
304///
305/// The adapters are only lent: each keeps the instance it came from alive, and that instance
306/// has to be dropped here, under the [`INSTANCES`] lock.
307#[cfg(not(target_family = "wasm"))]
308fn probe_with<R>(backend: wgpu::Backend, inspect: impl FnOnce(&[wgpu::Adapter]) -> R) -> R {
309    let _instances = INSTANCES.lock();
310    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
311        backends: backend.into(),
312        ..wgpu::InstanceDescriptor::new_without_display_handle()
313    });
314    inspect(&future::block_on(
315        instance.enumerate_adapters(backend.into()),
316    ))
317}
318
319/// The graphics API a device on `backend` comes up on.
320///
321/// The one pinned, where one is. Otherwise the first of the chain
322/// this machine has an adapter for — which is what makes `Auto` mean Vulkan
323/// wherever Vulkan exists, and the next thing where it does not.
324pub(crate) fn resolve_backend(backend: WgpuBackend) -> wgpu::Backend {
325    #[cfg(not(target_family = "wasm"))]
326    if let Some((api, _)) = settle(backend) {
327        return api;
328    }
329
330    // Nothing answered: hand back the first anyway, so the failure is the
331    // setup's own rather than a silent fallback to some other API.
332    backend_candidates(backend)[0]
333}
334
335/// The `DeviceId` addressing each adapter, in enumeration order.
336///
337/// Every device type counts from zero on its own: `WgpuDevice::DiscreteGpu(n)`
338/// is the nth *discrete* adapter, not the nth adapter overall, so a single
339/// counter over the mixed list hands out ids for devices that do not exist.
340/// `Cpu` carries no index in `WgpuDevice`, so it stays at zero.
341#[cfg(not(target_family = "wasm"))]
342fn adapter_device_ids(adapters: Vec<wgpu::AdapterInfo>) -> Vec<DeviceId> {
343    let mut next = [0u16; 7];
344
345    adapters
346        .into_iter()
347        .map(|adapter| {
348            let type_id = match adapter.device_type {
349                wgpu::DeviceType::DiscreteGpu => 0,
350                wgpu::DeviceType::IntegratedGpu => 1,
351                wgpu::DeviceType::VirtualGpu => 2,
352                wgpu::DeviceType::Cpu => 3,
353                wgpu::DeviceType::Other => 6,
354            };
355
356            // Only the indexed kinds have a counter; the rest are always zero.
357            let index = match next.get_mut(type_id as usize).filter(|_| type_id != 3) {
358                Some(next) => {
359                    let index = *next;
360                    *next += 1;
361                    index
362                }
363                None => 0,
364            };
365
366            DeviceId::new(type_id, index)
367        })
368        .collect()
369}
370
371/// A recoverable failure while acquiring or registering a wgpu runtime.
372#[derive(Debug, thiserror::Error)]
373#[non_exhaustive]
374pub enum WgpuInitError {
375    /// Options, environment defaults, or supplied setup handles are invalid.
376    #[error("Invalid wgpu configuration: {message}")]
377    InvalidConfiguration {
378        /// The invalid configuration and why it was rejected.
379        message: String,
380    },
381    /// The requested API is not compiled for this platform.
382    #[error("Graphics API {api:?} is unavailable in this build")]
383    UnsupportedGraphicsApi {
384        /// The requested API.
385        api: WgpuBackend,
386    },
387    /// No adapter matches the requested selector.
388    #[error("No adapter available for {device:?}: {message}")]
389    AdapterUnavailable {
390        /// The requested selector, including its graphics API.
391        device: WgpuDevice,
392        /// Details from adapter selection.
393        message: String,
394    },
395    /// The adapter refused to create a device.
396    #[error("Unable to request a wgpu device: {message}")]
397    RequestDevice {
398        /// Details from the graphics API.
399        message: String,
400    },
401    /// A runtime could not be registered.
402    #[error("Unable to register wgpu runtime: {message}")]
403    Registration {
404        /// Registration failure details.
405        message: String,
406    },
407}
408
409/// The values that control how a WGPU Runtime will perform its calculations.
410pub struct RuntimeOptions {
411    /// Control the amount of compute tasks to be aggregated into a single GPU command.
412    pub tasks_max: usize,
413    /// Configures the memory management.
414    pub memory_config: MemoryConfiguration,
415}
416
417impl Default for RuntimeOptions {
418    fn default() -> Self {
419        Self::try_default().expect("Invalid wgpu runtime defaults")
420    }
421}
422
423impl RuntimeOptions {
424    /// Resolve environment defaults without panicking on invalid configuration.
425    pub fn try_default() -> Result<Self, WgpuInitError> {
426        Self::with_tasks_max(None)
427    }
428
429    /// Resolve defaults, letting an explicit task count override the environment.
430    pub fn with_tasks_max(tasks_max: Option<usize>) -> Result<Self, WgpuInitError> {
431        let tasks_max = match tasks_max {
432            Some(value) => value,
433            None => match std::env::var("CUBECL_WGPU_MAX_TASKS") {
434                Ok(value) => value
435                    .parse()
436                    .map_err(|_| WgpuInitError::InvalidConfiguration {
437                        message: "CUBECL_WGPU_MAX_TASKS must be a positive integer".into(),
438                    })?,
439                Err(std::env::VarError::NotPresent) => 32,
440                Err(err) => {
441                    return Err(WgpuInitError::InvalidConfiguration {
442                        message: err.to_string(),
443                    });
444                }
445            },
446        };
447        let options = Self {
448            tasks_max,
449            memory_config: MemoryConfiguration::default(),
450        };
451        options.validate()?;
452        Ok(options)
453    }
454
455    fn validate(&self) -> Result<(), WgpuInitError> {
456        if self.tasks_max == 0 {
457            return Err(WgpuInitError::InvalidConfiguration {
458                message: "tasks_max must be greater than zero".into(),
459            });
460        }
461        Ok(())
462    }
463}
464
465/// A complete setup used to run wgpu.
466///
467/// These can either be created with [`init_setup`] or [`init_setup_async`].
468#[derive(Clone, Debug)]
469pub struct WgpuSetup {
470    /// The underlying wgpu instance.
471    pub instance: wgpu::Instance,
472    /// The selected 'adapter'. This corresponds to a physical device.
473    pub adapter: wgpu::Adapter,
474    /// The logical wgpu device the runtime will use.
475    pub device: wgpu::Device,
476    /// The queue Burn commands will be submitted to.
477    pub queue: wgpu::Queue,
478    /// The backend used by the setup.
479    pub backend: wgpu::Backend,
480}
481
482/// Create a [`WgpuDevice`] on an existing [`WgpuSetup`].
483/// Useful when you want to share a device between `CubeCL` and other wgpu-dependent libraries.
484///
485/// Uses [`AutoCompiler`] with WGSL fallback. Use [`init_device_with_api`] to request a graphics
486/// API explicitly, including native MSL on Metal when `msl` is enabled.
487///
488/// # Panics
489///
490/// If configuration is invalid, the setup is incompatible, or registration fails.
491/// Use [`try_init_device`] to handle initialization errors. Register each setup once,
492/// then clone the returned device to share its runtime; repeated handles are not detected.
493pub fn init_device(setup: WgpuSetup, options: RuntimeOptions) -> WgpuDevice {
494    try_init_device(setup, options).expect("Unable to register existing wgpu setup")
495}
496
497/// Import a setup with an explicitly selected graphics API.
498///
499/// With `msl` enabled, importing through [`crate::Metal`] requires native MSL support.
500/// [`AutoGraphicsApi`] permits WGSL fallback, as [`init_device`] does.
501///
502/// Like [`init_device`], each call generates a unique device ID; do not import the same setup twice.
503///
504/// # Panics
505///
506/// If configuration is invalid, the requested API differs from the setup's backend,
507/// the selected compiler is unavailable, or registration fails. Use
508/// [`try_init_device_with_api`] to handle these errors.
509pub fn init_device_with_api<G: GraphicsApi>(
510    setup: WgpuSetup,
511    options: RuntimeOptions,
512) -> WgpuDevice {
513    try_init_device_with_api::<G>(setup, options).expect("Unable to register existing wgpu setup")
514}
515
516/// Like [`init_setup_async`], but synchronous.
517/// On wasm, it is necessary to use [`init_setup_async`] instead.
518///
519/// A device brought up on a `G` other than [`AutoGraphicsApi`] is reached
520/// through the client this initializes, and is not among the devices
521/// [`Runtime::enumerate_devices`] lists: those ids index the auto backend's
522/// adapters, which is what a client created lazily resolves them against.
523///
524/// # Panics
525///
526/// Where `device` pins a graphics API and `G` names another: see
527/// [`init_setup_async`].
528/// With `msl` enabled, an explicitly selected Metal API also panics if native MSL is unavailable.
529pub fn init_setup<G: GraphicsApi>(device: &WgpuDevice, options: RuntimeOptions) -> WgpuSetup {
530    cfg_if::cfg_if! {
531        if #[cfg(target_family = "wasm")] {
532            let _ = (device, options);
533            panic!("Creating a wgpu setup synchronously is unsupported on wasm. Use init_async instead");
534        } else {
535            future::block_on(init_setup_async::<G>(device, options))
536        }
537    }
538}
539
540/// Initialize a client on the given device with the given options.
541/// This function is useful to configure the runtime options
542/// or to pick a different graphics API.
543///
544/// A device pinned to a graphics API comes up on that API: through
545/// [`AutoGraphicsApi`], or through the `G` naming the same one.
546///
547/// # Panics
548///
549/// Where `device` pins a graphics API and `G` names another. The client is
550/// registered under the device's id, and a pinned id promises its API to
551/// every caller who reaches for that client afterwards.
552/// With `msl` enabled, an explicitly selected Metal API also panics if native MSL is unavailable.
553pub async fn init_setup_async<G: GraphicsApi>(
554    device: &WgpuDevice,
555    options: RuntimeOptions,
556) -> WgpuSetup {
557    try_init_setup_async::<G>(device, options)
558        .await
559        .expect("Unable to initialize wgpu setup")
560}
561
562/// Register an existing setup without changing its enabled GPU capabilities.
563///
564/// Uses [`AutoCompiler`] with WGSL fallback. Use [`try_init_device_with_api`] to require
565/// an explicit graphics API and its compiler support. The supplied device must enable
566/// the features required by the selected compiler, including native extensions when
567/// using SPIR-V or MSL. Registration does not
568/// validate every native requirement. Each call allocates a new runtime identity;
569/// repeated registration of the same handles is not detected. Register once, then
570/// clone the returned device.
571pub fn try_init_device(
572    setup: WgpuSetup,
573    options: RuntimeOptions,
574) -> Result<WgpuDevice, WgpuInitError> {
575    try_init_device_with_api::<AutoGraphicsApi>(setup, options)
576}
577
578/// Register an existing setup with an explicitly selected graphics API.
579///
580/// The requested API must match the setup's backend. With `msl` enabled, an explicit
581/// [`crate::Metal`] request requires native MSL support; [`AutoGraphicsApi`] permits
582/// WGSL fallback. Invalid configuration, unavailable compilers, and registration
583/// failures are returned as errors. Register once, then clone the returned device.
584pub fn try_init_device_with_api<G: GraphicsApi>(
585    setup: WgpuSetup,
586    options: RuntimeOptions,
587) -> Result<WgpuDevice, WgpuInitError> {
588    use core::sync::atomic::{AtomicU32, Ordering};
589    static COUNTER: AtomicU32 = AtomicU32::new(0);
590    options.validate()?;
591    let requested_backend = G::backend_for(&WgpuDevice::default());
592    if requested_backend != WgpuBackend::Auto
593        && !backend_candidates(requested_backend).contains(&setup.backend)
594    {
595        return Err(WgpuInitError::InvalidConfiguration {
596            message: format!(
597                "the imported setup uses {:?}, but {requested_backend:?} was requested",
598                setup.backend
599            ),
600        });
601    }
602    if setup.backend != setup.adapter.get_info().backend
603        || setup.backend != setup.device.adapter_info().backend
604    {
605        return Err(WgpuInitError::InvalidConfiguration {
606            message: "setup backend does not match its adapter and device".into(),
607        });
608    }
609    let limits = setup.device.limits();
610    if limits.max_compute_invocations_per_workgroup == 0
611        || limits.max_storage_buffers_per_shader_stage < 2
612    {
613        return Err(WgpuInitError::InvalidConfiguration {
614            message: "compute support and at least two storage bindings are required".into(),
615        });
616    }
617    // External DeviceId indices are u16. Never wrap an ID onto a live runtime.
618    let id = COUNTER
619        .try_update(Ordering::Relaxed, Ordering::Relaxed, |id| {
620            (id <= u16::MAX as u32).then_some(id + 1)
621        })
622        .map_err(|_| WgpuInitError::Registration {
623            message: "external wgpu device IDs exhausted".into(),
624        })?;
625    let device = WgpuDevice::new(WgpuDeviceKind::Existing(id));
626    let server =
627        try_create_server::<AutoCompiler>(setup, options, device.to_id(), requested_backend)?;
628    Client::try_init(device.to_id(), server).map_err(|err| WgpuInitError::Registration {
629        message: format!("{device:?}: {err:?}"),
630    })?;
631    Ok(device)
632}
633
634/// Fallible synchronous initialization. Use [`try_init_setup_async`] on wasm.
635#[cfg(not(target_family = "wasm"))]
636pub fn try_init_setup<G: GraphicsApi>(
637    device: &WgpuDevice,
638    options: RuntimeOptions,
639) -> Result<WgpuSetup, WgpuInitError> {
640    future::block_on(try_init_setup_async::<G>(device, options))
641}
642
643/// Initialize the selected runtime, returning acquisition or registration failures.
644///
645/// An initialized selector cannot be configured again. Registration is atomic in
646/// `CubeCL`'s service registry, but concurrent calls may each acquire a device before
647/// one registers successfully. Failed or cancelled acquisition leaves no registration.
648/// With `msl` enabled, explicitly selecting Metal requires native MSL support;
649/// automatic selection permits WGSL fallback. Compiler validation failures are returned
650/// as [`WgpuInitError::InvalidConfiguration`] before registration.
651pub async fn try_init_setup_async<G: GraphicsApi>(
652    device: &WgpuDevice,
653    options: RuntimeOptions,
654) -> Result<WgpuSetup, WgpuInitError> {
655    options.validate()?;
656    if matches!(device.kind, WgpuDeviceKind::Existing(_)) {
657        return Err(WgpuInitError::InvalidConfiguration {
658            message:
659                "an Existing ID is not an adapter selector; register a setup or reuse its device"
660                    .into(),
661        });
662    }
663    // Preserve Auto until acquisition resolves it, so automatic Metal selection
664    // can fall back to WGSL while an explicit Metal request requires native MSL.
665    let requested_backend = G::backend_for(device);
666    if device.backend != WgpuBackend::Auto && device.backend != requested_backend {
667        return Err(WgpuInitError::InvalidConfiguration {
668            message: format!("{device:?} cannot be initialized on {requested_backend:?}"),
669        });
670    }
671    let mut selected = device.clone();
672    selected.backend = requested_backend;
673    let setup = try_create_setup(&selected).await?;
674    let server = try_create_server::<AutoCompiler>(
675        setup.clone(),
676        options,
677        device.to_id(),
678        requested_backend,
679    )?;
680    Client::try_init(device.to_id(), server).map_err(|err| WgpuInitError::Registration {
681        message: format!("{device:?}: {err:?}"),
682    })?;
683    Ok(setup)
684}
685
686/// The runtime name for `backend`, naming the compiler that serves it.
687///
688/// Read from the same options [`AutoCompiler`] picks its compiler from, so a device that fell
689/// back to WGSL says so, along with its API: DX12 and Vulkan on one card share a fingerprint.
690fn runtime_name(backend: wgpu::Backend, options: &WgpuCompilationOptions) -> &'static str {
691    match backend {
692        wgpu::Backend::Vulkan if options.supports_vulkan_compiler => "wgpu<spirv>",
693        wgpu::Backend::Metal if options.supports_msl_compiler => "wgpu<msl>",
694        wgpu::Backend::Vulkan => "wgpu<wgsl-vulkan>",
695        wgpu::Backend::Metal => "wgpu<wgsl-metal>",
696        wgpu::Backend::Dx12 => "wgpu<wgsl-dx12>",
697        wgpu::Backend::Gl => "wgpu<wgsl-gl>",
698        _ => "wgpu<wgsl>",
699    }
700}
701
702fn try_create_server<C: WgpuCompiler>(
703    setup: WgpuSetup,
704    options: RuntimeOptions,
705    device_id: DeviceId,
706    requested_backend: WgpuBackend,
707) -> Result<WgpuServer<C>, WgpuInitError> {
708    let limits = setup.device.limits();
709    let adapter_limits = setup.adapter.limits();
710    let mut adapter_info = setup.adapter.get_info();
711
712    // Workaround: WebGPU reports some "fake" subgroup info atm, as it's not really supported yet.
713    // However, some algorithms do rely on having this information eg. cubecl-reduce uses max subgroup size _even_ when
714    // subgroups aren't used. For now, just override with the maximum range of subgroups possible.
715    if adapter_info.subgroup_min_size == 0 && adapter_info.subgroup_max_size == 0 {
716        // There is in theory nothing limiting the size to go below 8 but in practice 8 is the minimum found anywhere.
717        adapter_info.subgroup_min_size = 8;
718        // This is a hard limit of GPU APIs (subgroup ballot returns 4 * 32 bits).
719        adapter_info.subgroup_max_size = 128;
720    }
721
722    // WebGPU states no capacity. `register_features` fills it in where the
723    // backend's own API states one and this build enabled that backend:
724    // `spirv` for Vulkan, `msl` for Metal.
725    let mem_props = MemoryDeviceProperties::new(
726        limits
727            .max_storage_buffer_binding_size
728            .min(limits.max_buffer_size),
729        limits
730            .min_uniform_buffer_offset_alignment
731            .max(limits.min_storage_buffer_offset_alignment) as u64,
732    );
733    let max_count = adapter_limits.max_compute_workgroups_per_dimension;
734    let hardware_props = HardwareProperties {
735        load_width: 128,
736        vector_register_count: None,
737        // On Apple Silicon, the plane size is 32,
738        // though the minimum and maximum differ.
739        // https://github.com/gpuweb/gpuweb/issues/3950
740        #[cfg(apple_silicon)]
741        plane_size_min: 32,
742        #[cfg(not(apple_silicon))]
743        plane_size_min: adapter_info.subgroup_min_size,
744        #[cfg(apple_silicon)]
745        plane_size_max: 32,
746        #[cfg(not(apple_silicon))]
747        plane_size_max: adapter_info.subgroup_max_size,
748        // wgpu uses an additional buffer for variable-length buffers,
749        // so we have to use one buffer less on our side to make room for that wgpu internal buffer.
750        // See: https://github.com/gfx-rs/wgpu/blob/a9638c8e3ac09ce4f27ac171f8175671e30365fd/wgpu-hal/src/metal/device.rs#L799
751        max_bindings: limits
752            .max_storage_buffers_per_shader_stage
753            .saturating_sub(1),
754        max_shared_memory_size: limits.max_compute_workgroup_storage_size as usize,
755        max_cube_count: (max_count, max_count, max_count),
756        max_units_per_cube: adapter_limits.max_compute_invocations_per_workgroup,
757        max_cube_dim: (
758            adapter_limits.max_compute_workgroup_size_x,
759            adapter_limits.max_compute_workgroup_size_y,
760            adapter_limits.max_compute_workgroup_size_z,
761        ),
762        num_streaming_multiprocessors: None,
763        num_tensor_cores: None,
764        min_tensor_cores_dim: None,
765        num_cpu_cores: None, // TODO: Check if device is CPU.
766        last_level_cache_size: None,
767        max_vector_size: 4,
768        // Init later if extension is enabled
769        cube_mma_reserved_shared_memory: 0,
770    };
771
772    let mut compilation_options = Default::default();
773
774    let features = setup.adapter.features();
775
776    let time_measurement = if features.contains(wgpu::Features::TIMESTAMP_QUERY) {
777        TimingMethod::Device
778    } else {
779        TimingMethod::System
780    };
781
782    // The adapter's vendor/device pair, which is what `WgpuServer` keys its
783    // SPIR-V store on. Reported unconditionally, even in a WGSL-only build that
784    // persists no compiled code: the measurement caches (autotune, throughput)
785    // are named by it beside the adapter name, which alone is not an identity.
786    let fingerprint = format!("spirv_{}_{}", adapter_info.vendor, adapter_info.device);
787
788    let mut device_props = DeviceProperties::new(
789        Default::default(),
790        mem_props,
791        hardware_props,
792        time_measurement,
793        DeviceIdentity {
794            name: adapter_info.name.clone(),
795            fingerprint,
796            physical: backend::physical_device(&setup.adapter, &adapter_info),
797        },
798    );
799
800    // Metal reports no core count; the GPU's driver states it in the I/O Registry entry the card
801    // is named by.
802    #[cfg(target_os = "macos")]
803    if let Some(entry) =
804        (device_props.identity.physical.as_ref()).and_then(|card| card.registry_entry_id)
805    {
806        device_props.hardware.num_streaming_multiprocessors =
807            cubecl_monitoring::gpu_utilization::gpu_core_count(entry);
808    }
809
810    #[cfg(not(all(target_os = "macos", feature = "msl")))]
811    {
812        if features.contains(wgpu::Features::SUBGROUP)
813            && setup.adapter.get_info().device_type != wgpu::DeviceType::Cpu
814        {
815            use cubecl_ir::features::Plane;
816
817            device_props.features.plane.insert(Plane::Ops);
818        }
819    }
820
821    backend::register_features(&setup.adapter, &mut device_props, &mut compilation_options);
822
823    // PATCH: remove once wgpu supports `VK_EXT_subgroup_size_control` and this runtime requires
824    // full subgroups, so the adapter commits to an exact width on its own.
825    //
826    // A 32 to 64 plane range with plane ops is AMD RDNA under Vulkan: the driver may compile a
827    // shader wave32 or wave64, so the range is honest and a kernel needing one width rightly
828    // refuses it. On an AMD integrated GPU, though, RADV dispatches these compute shaders wave64
829    // (validated on gfx1151, where a wave32 dispatch would halve every plane sum), so the range
830    // narrows to 64 there. Only on a machine whose GPU is that integrated one, the only machines
831    // it was validated on: discrete RDNA keeps the range.
832    #[cfg(not(target_family = "wasm"))]
833    if device_props
834        .features
835        .plane
836        .contains(cubecl_ir::features::Plane::Ops)
837        && device_props.hardware.plane_size_min == 32
838        && device_props.hardware.plane_size_max == 64
839        && adapter_info.device_type == wgpu::DeviceType::IntegratedGpu
840        && !probe(setup.backend)
841            .iter()
842            .any(|adapter| adapter.device_type == wgpu::DeviceType::DiscreteGpu)
843    {
844        device_props.hardware.plane_size_min = 64;
845    }
846
847    #[cfg(any(feature = "spirv", feature = "msl"))]
848    if compilation_options.supports_vulkan_compiler || compilation_options.supports_msl_compiler {
849        device_props
850            .features
851            .plane
852            .insert(cubecl_ir::features::Plane::NonUniformControlFlow);
853    }
854
855    C::validate_runtime(setup.backend, &compilation_options, requested_backend).map_err(|err| {
856        WgpuInitError::InvalidConfiguration {
857            message: format!("failed to initialize the wgpu compiler: {err}"),
858        }
859    })?;
860
861    let compiler = C::init(setup.backend, &compilation_options);
862    compiler.limit_properties(&mut device_props);
863
864    let logger = alloc::sync::Arc::new(ServerLogger::default());
865    let name = runtime_name(setup.backend, &compilation_options);
866
867    let allocator = ContiguousMemoryLayoutPolicy::new(device_props.memory.alignment as usize);
868    let memory_properties = device_props.memory.clone();
869    let (utilities, captures) = ServerUtilities::init(
870        ServiceId::of::<WgpuServer<C>>(device_id),
871        name,
872        device_props,
873        compiler.target_properties(),
874        logger,
875        allocator,
876    );
877    Ok(WgpuServer::new(
878        memory_properties,
879        options.memory_config,
880        compilation_options,
881        setup.device.clone(),
882        setup.queue,
883        options.tasks_max,
884        setup.backend,
885        time_measurement,
886        utilities,
887        captures,
888    ))
889}
890
891/// Acquire an adapter and device without blocking the caller's executor.
892async fn try_create_setup(device: &WgpuDevice) -> Result<WgpuSetup, WgpuInitError> {
893    // Validate overrides before touching the graphics APIs. They choose hardware, not
894    // the graphics API or the registration key the caller selected.
895    let override_device = match device.kind {
896        WgpuDeviceKind::DefaultDevice => try_device_override()?.map(|kind| WgpuDevice {
897            kind,
898            backend: device.backend,
899        }),
900        _ => None,
901    };
902    let device = override_device.as_ref().unwrap_or(device);
903    match device.kind {
904        WgpuDeviceKind::DiscreteGpu(index)
905        | WgpuDeviceKind::IntegratedGpu(index)
906        | WgpuDeviceKind::VirtualGpu(index)
907        | WgpuDeviceKind::Other(index)
908            if index > WgpuDeviceKind::MAX_INDEX =>
909        {
910            return Err(WgpuInitError::InvalidConfiguration {
911                message: format!(
912                    "adapter index {index} exceeds {}",
913                    WgpuDeviceKind::MAX_INDEX
914                ),
915            });
916        }
917        WgpuDeviceKind::Existing(_) => {
918            return Err(WgpuInitError::InvalidConfiguration {
919                message: "an Existing ID cannot acquire an adapter".into(),
920            });
921        }
922        _ => {}
923    }
924    let candidates = backend_candidates(device.backend);
925    if !candidates
926        .iter()
927        .any(|api| wgpu::Instance::enabled_backend_features().contains((*api).into()))
928    {
929        return Err(WgpuInitError::UnsupportedGraphicsApi {
930            api: device.backend,
931        });
932    }
933    let mut backend = None;
934    #[cfg(not(target_family = "wasm"))]
935    for api in candidates {
936        if !wgpu::Instance::enabled_backend_features().contains(api.into()) {
937            continue;
938        }
939        let adapters = probe(api);
940        if adapters
941            .iter()
942            .any(|adapter| adapter.device_type != wgpu::DeviceType::Cpu)
943        {
944            backend = Some(api);
945            break;
946        }
947        if backend.is_none() && !adapters.is_empty() {
948            backend = Some(api);
949        }
950    }
951    #[cfg(target_family = "wasm")]
952    for api in candidates {
953        if api == wgpu::Backend::BrowserWebGpu
954            && wgpu::Instance::enabled_backend_features().contains(api.into())
955        {
956            backend = Some(api);
957            break;
958        }
959    }
960    let backend = backend.ok_or_else(|| WgpuInitError::AdapterUnavailable {
961        device: device.clone(),
962        message: "no adapter found on the requested graphics API".into(),
963    })?;
964    let (instance, adapter) = request_adapter(device, backend).await?;
965    let (device, queue) = backend::try_request_device(&adapter).await?;
966    Ok(WgpuSetup {
967        instance,
968        adapter,
969        device,
970        queue,
971        backend,
972    })
973}
974
975async fn request_adapter(
976    device: &WgpuDevice,
977    backend: wgpu::Backend,
978) -> Result<(wgpu::Instance, wgpu::Adapter), WgpuInitError> {
979    #[cfg(not(feature = "vulkan-validate"))]
980    let instance_flags = {
981        let debug = ServerLogger::default();
982        // Debug/validation layers cost real per-dispatch CPU time, so only
983        // source-level compilation logging (`full`) opts into them — `basic`
984        // is passive name-only logging and must not change how kernels run.
985        match (debug.profile_level(), debug.compilation_source_activated()) {
986            (Some(ProfileLevel::Full), _) => InstanceFlags::advanced_debugging(),
987            (_, true) => InstanceFlags::debugging(),
988            (_, false) => InstanceFlags::default(),
989        }
990    };
991    #[cfg(feature = "vulkan-validate")]
992    let instance_flags = InstanceFlags::advanced_debugging();
993    log::debug!("{instance_flags:?}");
994    let instance = device_instance(backend, instance_flags);
995
996    let adapter = match device.kind {
997        #[cfg(not(target_family = "wasm"))]
998        WgpuDeviceKind::DiscreteGpu(num) => {
999            select_from_adapter_list(
1000                num,
1001                "No Discrete GPU device found",
1002                &instance,
1003                device,
1004                backend,
1005            )
1006            .await
1007        }
1008        #[cfg(not(target_family = "wasm"))]
1009        WgpuDeviceKind::IntegratedGpu(num) => {
1010            select_from_adapter_list(
1011                num,
1012                "No Integrated GPU device found",
1013                &instance,
1014                device,
1015                backend,
1016            )
1017            .await
1018        }
1019        #[cfg(not(target_family = "wasm"))]
1020        WgpuDeviceKind::VirtualGpu(num) => {
1021            select_from_adapter_list(
1022                num,
1023                "No Virtual GPU device found",
1024                &instance,
1025                device,
1026                backend,
1027            )
1028            .await
1029        }
1030        #[cfg(not(target_family = "wasm"))]
1031        WgpuDeviceKind::Other(num) => {
1032            select_from_adapter_list(num, "No Other device found", &instance, device, backend).await
1033        }
1034        #[cfg(not(target_family = "wasm"))]
1035        WgpuDeviceKind::Cpu => {
1036            select_from_adapter_list(0, "No CPU device found", &instance, device, backend).await
1037        }
1038        #[cfg(target_family = "wasm")]
1039        WgpuDeviceKind::IntegratedGpu(_) => {
1040            request_adapter_with_preference(&instance, wgpu::PowerPreference::LowPower).await
1041        }
1042        WgpuDeviceKind::Existing(_) => {
1043            unreachable!("Cannot select an adapter for an existing device.")
1044        }
1045        _ => {
1046            request_adapter_with_preference(&instance, wgpu::PowerPreference::HighPerformance).await
1047        }
1048    }
1049    .map_err(|message| WgpuInitError::AdapterUnavailable {
1050        device: device.clone(),
1051        message,
1052    })?;
1053
1054    Ok((instance, adapter))
1055}
1056
1057/// The instance devices on `backend` come up on, created with the first device's `flags`.
1058fn device_instance(backend: wgpu::Backend, flags: InstanceFlags) -> wgpu::Instance {
1059    let mut instances = INSTANCES.lock();
1060    if let Some((_, instance)) = instances.iter().find(|(api, _)| *api == backend) {
1061        return instance.clone();
1062    }
1063    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
1064        backends: backend.into(),
1065        flags,
1066        ..wgpu::InstanceDescriptor::new_without_display_handle()
1067    });
1068    instances.push((backend, instance.clone()));
1069    instance
1070}
1071
1072async fn request_adapter_with_preference(
1073    instance: &wgpu::Instance,
1074    power_preference: wgpu::PowerPreference,
1075) -> Result<wgpu::Adapter, String> {
1076    instance
1077        .request_adapter(&RequestAdapterOptions {
1078            power_preference,
1079            force_fallback_adapter: false,
1080            compatible_surface: None,
1081            ..RequestAdapterOptions::default()
1082        })
1083        .await
1084        .map_err(|err| err.to_string())
1085}
1086
1087#[cfg(not(target_family = "wasm"))]
1088async fn select_from_adapter_list(
1089    num: usize,
1090    error: &str,
1091    instance: &wgpu::Instance,
1092    device: &WgpuDevice,
1093    backend: wgpu::Backend,
1094) -> Result<wgpu::Adapter, String> {
1095    // A kind is what the graphics API reports, and nothing stands in for it:
1096    // `OpenGL` calling a GPU `Other` makes it `Other(n)` there, not a discrete
1097    // GPU by another name. Anything looser selects adapters `find_device` says
1098    // the machine does not have, and two ids end up on one adapter.
1099    let adapters = instance.enumerate_adapters(backend.into()).await;
1100    let found = adapters
1101        .iter()
1102        .map(|adapter| adapter.get_info())
1103        .collect::<Vec<_>>();
1104
1105    let is_same_type = |adapter: &wgpu::Adapter| {
1106        let device_type = adapter.get_info().device_type;
1107
1108        match device.kind {
1109            WgpuDeviceKind::DiscreteGpu(_) => device_type == wgpu::DeviceType::DiscreteGpu,
1110            WgpuDeviceKind::IntegratedGpu(_) => device_type == wgpu::DeviceType::IntegratedGpu,
1111            WgpuDeviceKind::VirtualGpu(_) => device_type == wgpu::DeviceType::VirtualGpu,
1112            WgpuDeviceKind::Cpu => device_type == wgpu::DeviceType::Cpu,
1113            WgpuDeviceKind::Other(_) => device_type == wgpu::DeviceType::Other,
1114            WgpuDeviceKind::DefaultDevice => true,
1115            WgpuDeviceKind::Existing(_) => {
1116                unreachable!("Cannot select an adapter for an existing device.")
1117            }
1118        }
1119    };
1120
1121    adapters
1122        .into_iter()
1123        .filter(is_same_type)
1124        .nth(num)
1125        .ok_or_else(|| format!("{error}, adapters {found:?}"))
1126}
1127
1128fn try_device_override() -> Result<Option<WgpuDeviceKind>, WgpuInitError> {
1129    let value = match std::env::var("CUBECL_WGPU_DEFAULT_DEVICE") {
1130        Ok(value) => value,
1131        Err(std::env::VarError::NotPresent) => return Ok(None),
1132        Err(err) => {
1133            return Err(WgpuInitError::InvalidConfiguration {
1134                message: err.to_string(),
1135            });
1136        }
1137    };
1138    let parsed = if value == "Cpu" {
1139        Some(WgpuDeviceKind::Cpu)
1140    } else {
1141        [
1142            (
1143                "DiscreteGpu(",
1144                WgpuDeviceKind::DiscreteGpu as fn(usize) -> WgpuDeviceKind,
1145            ),
1146            ("IntegratedGpu(", WgpuDeviceKind::IntegratedGpu),
1147            ("VirtualGpu(", WgpuDeviceKind::VirtualGpu),
1148            ("Other(", WgpuDeviceKind::Other),
1149        ]
1150        .into_iter()
1151        .find_map(|(prefix, kind)| {
1152            value
1153                .strip_prefix(prefix)?
1154                .strip_suffix(')')?
1155                .parse()
1156                .ok()
1157                .map(kind)
1158        })
1159    };
1160    parsed
1161        .map(Some)
1162        .ok_or_else(|| WgpuInitError::InvalidConfiguration {
1163            message: format!("Invalid CUBECL_WGPU_DEFAULT_DEVICE: {value}"),
1164        })
1165}
1166
1167#[cfg(all(test, not(target_family = "wasm")))]
1168mod device_tests {
1169    use super::*;
1170
1171    #[test]
1172    #[ignore = "requires a native Vulkan adapter (software adapters are supported)"]
1173    fn fallible_import_preserves_api_selection_and_unique_ids() {
1174        let selector = WgpuDevice::default().on(WgpuBackend::Vulkan);
1175        let setup = future::block_on(try_create_setup(&selector)).unwrap();
1176
1177        assert!(matches!(
1178            try_init_device_with_api::<crate::Metal>(setup.clone(), RuntimeOptions::default()),
1179            Err(WgpuInitError::InvalidConfiguration { .. })
1180        ));
1181
1182        let automatic = try_init_device(setup, RuntimeOptions::default()).unwrap();
1183        let setup = future::block_on(try_create_setup(&selector)).unwrap();
1184        let explicit = init_device_with_api::<crate::Vulkan>(setup, RuntimeOptions::default());
1185        assert_ne!(automatic.to_id(), explicit.to_id());
1186    }
1187
1188    #[cfg(feature = "msl")]
1189    #[test]
1190    #[ignore = "requires a native Vulkan adapter (software adapters are supported)"]
1191    fn compiler_validation_returns_an_initialization_error() {
1192        let selector = WgpuDevice::default().on(WgpuBackend::Vulkan);
1193        let setup = future::block_on(try_create_setup(&selector)).unwrap();
1194
1195        // An explicit MSL compiler still rejects Vulkan, through a recoverable error.
1196        assert!(matches!(
1197            try_create_server::<crate::MslCompiler>(
1198                setup.clone(),
1199                RuntimeOptions::default(),
1200                selector.to_id(),
1201                WgpuBackend::Auto,
1202            ),
1203            Err(WgpuInitError::InvalidConfiguration { message })
1204                if message.contains("MslCompiler") && message.contains("Metal")
1205        ));
1206
1207        // AutoCompiler must also receive the original API request for validation.
1208        assert!(matches!(
1209            try_create_server::<AutoCompiler>(
1210                setup,
1211                RuntimeOptions::default(),
1212                selector.to_id(),
1213                WgpuBackend::Metal,
1214            ),
1215            Err(WgpuInitError::InvalidConfiguration { .. })
1216        ));
1217
1218        // Failed validation left no registration behind; a supported request succeeds.
1219        let setup =
1220            try_init_setup::<AutoGraphicsApi>(&selector, RuntimeOptions::default()).unwrap();
1221        assert_eq!(setup.backend, wgpu::Backend::Vulkan);
1222        assert!(matches!(
1223            try_init_setup::<AutoGraphicsApi>(&selector, RuntimeOptions::default()),
1224            Err(WgpuInitError::Registration { .. })
1225        ));
1226    }
1227
1228    const PINNED: [WgpuBackend; 4] = [
1229        WgpuBackend::Vulkan,
1230        WgpuBackend::Metal,
1231        WgpuBackend::Dx12,
1232        WgpuBackend::Gl,
1233    ];
1234
1235    /// One adapter is one device. Listing it again for every graphics API
1236    /// that reaches it makes a single GPU look like several, and whatever
1237    /// counts devices — a collective, a transfer between two of them — runs
1238    /// on hardware that is not there.
1239    #[test]
1240    fn each_adapter_is_listed_once() {
1241        let ids = <WgpuRuntime>::enumerate_all_devices();
1242
1243        let adapters = settle(WgpuBackend::Auto).map_or(0, |(_, adapters)| adapters.len());
1244
1245        assert_eq!(ids.len(), adapters);
1246        for id in ids {
1247            assert_eq!(WgpuDevice::from_id(id).backend, WgpuBackend::Auto);
1248        }
1249    }
1250
1251    /// A device pinned to an API is found where that API has it, whatever
1252    /// the API `Auto` settles on has — the kinds differ from one API to the
1253    /// next, `OpenGL` calling a GPU what Vulkan calls discrete.
1254    #[test]
1255    fn a_device_is_found_on_the_api_it_names() {
1256        for backend in PINNED {
1257            let reachable = adapters_on(backend);
1258
1259            for id in reachable.iter() {
1260                assert_eq!(
1261                    <WgpuRuntime>::find_device(*id),
1262                    Ok(()),
1263                    "{id} on {backend:?}"
1264                );
1265            }
1266
1267            let default = WgpuDevice::new(WgpuDeviceKind::DefaultDevice).on(backend);
1268            assert_eq!(
1269                <WgpuRuntime>::find_device(default.to_id()).is_ok(),
1270                !reachable.is_empty(),
1271                "the default device on {backend:?}"
1272            );
1273        }
1274    }
1275
1276    /// The default device stands for whichever adapter it lands on, so it is
1277    /// its own only peer: the adapters listed beside it would give one of them
1278    /// a second client, and leaving it out drops the caller from its own list.
1279    #[test]
1280    fn the_default_device_is_its_own_only_peer() {
1281        for backend in PINNED.into_iter().chain([WgpuBackend::Auto]) {
1282            let default = WgpuDevice::new(WgpuDeviceKind::DefaultDevice)
1283                .on(backend)
1284                .to_id();
1285
1286            let expected = match adapters_on(backend).is_empty() {
1287                true => Vec::new(),
1288                false => alloc::vec![default],
1289            };
1290
1291            assert_eq!(<WgpuRuntime>::enumerate_devices_like(default), expected);
1292        }
1293    }
1294
1295    /// Setting a pinned device up through `AutoGraphicsApi` keeps its pin —
1296    /// the natural call, `init_setup` being the only way to pass options.
1297    #[test]
1298    fn auto_defers_to_the_api_a_device_pins() {
1299        for (backend, api) in [
1300            (WgpuBackend::Vulkan, wgpu::Backend::Vulkan),
1301            (WgpuBackend::Metal, wgpu::Backend::Metal),
1302            (WgpuBackend::Dx12, wgpu::Backend::Dx12),
1303            (WgpuBackend::Gl, wgpu::Backend::Gl),
1304        ] {
1305            let device = WgpuDevice::new(WgpuDeviceKind::DefaultDevice).on(backend);
1306
1307            assert_eq!(AutoGraphicsApi::backend_for(&device), backend);
1308            assert_eq!(resolve_backend(AutoGraphicsApi::backend_for(&device)), api);
1309        }
1310    }
1311
1312    /// Naming one API for a device pinned to another is refused before any
1313    /// adapter is asked for: the client lands under the pinned id, and would
1314    /// hand every later caller the wrong API.
1315    #[test]
1316    fn a_setup_on_another_api_than_the_pinned_one_is_refused() {
1317        let device = WgpuDevice::new(WgpuDeviceKind::DefaultDevice).on(WgpuBackend::Gl);
1318
1319        assert!(matches!(
1320            try_init_setup::<crate::Vulkan>(&device, RuntimeOptions::default()),
1321            Err(WgpuInitError::InvalidConfiguration { .. })
1322        ));
1323    }
1324
1325    /// A pinned device's peers are those of its own API. The same adapters on
1326    /// `Auto` are other devices, with other clients.
1327    #[test]
1328    fn a_pinned_device_is_enumerated_with_its_own_api() {
1329        for backend in PINNED {
1330            for id in adapters_on(backend) {
1331                let peers = <WgpuRuntime>::enumerate_devices_like(id);
1332
1333                assert!(peers.contains(&id), "{id} among {peers:?}");
1334                for peer in peers {
1335                    assert_eq!(WgpuDevice::from_id(peer).backend, backend);
1336                }
1337            }
1338        }
1339    }
1340
1341    /// And one that API does not have is a miss, reported against what it
1342    /// has of that kind.
1343    #[test]
1344    fn an_index_past_the_end_is_not_found_on_any_api() {
1345        for backend in PINNED.into_iter().chain([WgpuBackend::Auto]) {
1346            let device = WgpuDevice::new(WgpuDeviceKind::DiscreteGpu(4242)).on(backend);
1347
1348            let discrete = adapters_on(backend)
1349                .into_iter()
1350                .filter(|id| id.type_id == device.to_id().type_id)
1351                .count();
1352
1353            assert_eq!(<WgpuRuntime>::find_device(device.to_id()), Err(discrete));
1354        }
1355    }
1356
1357    #[test]
1358    fn wgsl_runtime_names_its_api() {
1359        let wgsl = WgpuCompilationOptions::default();
1360        let vulkan = runtime_name(wgpu::Backend::Vulkan, &wgsl);
1361        assert_ne!(vulkan, runtime_name(wgpu::Backend::Gl, &wgsl));
1362        assert_ne!(vulkan, runtime_name(wgpu::Backend::Dx12, &wgsl));
1363    }
1364}