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#[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 fn target_properties() -> TargetProperties {
85 TargetProperties::default()
86 }
87
88 fn enumerate_devices(type_id: u16) -> Vec<DeviceId> {
89 Self::enumerate_devices_like(DeviceId::new(type_id, 0))
92 }
93
94 fn is_available() -> bool {
95 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 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 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 #[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 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 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
197fn adapters_on(backend: WgpuBackend) -> Vec<DeviceId> {
204 #[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#[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
252fn 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
266static INSTANCES: Mutex<Vec<(wgpu::Backend, wgpu::Instance)>> = Mutex::new(Vec::new());
270
271#[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#[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 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#[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
319pub(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 backend_candidates(backend)[0]
333}
334
335#[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 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#[derive(Debug, thiserror::Error)]
373#[non_exhaustive]
374pub enum WgpuInitError {
375 #[error("Invalid wgpu configuration: {message}")]
377 InvalidConfiguration {
378 message: String,
380 },
381 #[error("Graphics API {api:?} is unavailable in this build")]
383 UnsupportedGraphicsApi {
384 api: WgpuBackend,
386 },
387 #[error("No adapter available for {device:?}: {message}")]
389 AdapterUnavailable {
390 device: WgpuDevice,
392 message: String,
394 },
395 #[error("Unable to request a wgpu device: {message}")]
397 RequestDevice {
398 message: String,
400 },
401 #[error("Unable to register wgpu runtime: {message}")]
403 Registration {
404 message: String,
406 },
407}
408
409pub struct RuntimeOptions {
411 pub tasks_max: usize,
413 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 pub fn try_default() -> Result<Self, WgpuInitError> {
426 Self::with_tasks_max(None)
427 }
428
429 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#[derive(Clone, Debug)]
469pub struct WgpuSetup {
470 pub instance: wgpu::Instance,
472 pub adapter: wgpu::Adapter,
474 pub device: wgpu::Device,
476 pub queue: wgpu::Queue,
478 pub backend: wgpu::Backend,
480}
481
482pub fn init_device(setup: WgpuSetup, options: RuntimeOptions) -> WgpuDevice {
494 try_init_device(setup, options).expect("Unable to register existing wgpu setup")
495}
496
497pub 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
516pub 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
540pub 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
562pub 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
578pub 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 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#[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
643pub 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 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
686fn 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 if adapter_info.subgroup_min_size == 0 && adapter_info.subgroup_max_size == 0 {
716 adapter_info.subgroup_min_size = 8;
718 adapter_info.subgroup_max_size = 128;
720 }
721
722 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 #[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 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, last_level_cache_size: None,
767 max_vector_size: 4,
768 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 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 #[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 #[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
891async fn try_create_setup(device: &WgpuDevice) -> Result<WgpuSetup, WgpuInitError> {
893 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 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
1057fn 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 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 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 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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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}