ruda_driver_wgpu/runtime/
device_service.rs

1use super::*;
2
3impl DeviceService for WgpuServer {
4    fn init(device_id: ruda_core::device::DeviceId) -> Self {
5        let device = WgpuDevice::from_id(device_id);
6        let setup = future::block_on(create_setup_for_device(&device, AutoGraphicsApi::backend()));
7        create_server(setup, RuntimeOptions::default())
8    }
9
10    fn utilities(&self) -> ServerUtilitiesHandle {
11        self.utilities.clone() as ServerUtilitiesHandle
12    }
13}
14
15pub(crate) fn create_server(setup: WgpuSetup, options: RuntimeOptions) -> WgpuServer {
16    let limits = setup.device.limits();
17    let adapter_limits = setup.adapter.limits();
18    let mut adapter_info = setup.adapter.get_info();
19
20    // Workaround: WebGPU reports some "fake" subgroup info atm, as it's not really supported yet.
21    // However, some algorithms do rely on having this information eg. ruPRIM uses max subgroup size _even_ when
22    // subgroups aren't used. For now, just override with the maximum range of subgroups possible.
23    if adapter_info.subgroup_min_size == 0 && adapter_info.subgroup_max_size == 0 {
24        // There is in theory nothing limiting the size to go below 8 but in practice 8 is the minimum found anywhere.
25        adapter_info.subgroup_min_size = 8;
26        // This is a hard limit of GPU APIs (subgroup ballot returns 4 * 32 bits).
27        adapter_info.subgroup_max_size = 128;
28    }
29
30    let mem_props = MemoryDeviceProperties {
31        max_page_size: limits.max_storage_buffer_binding_size,
32        alignment: limits.min_uniform_buffer_offset_alignment as u64,
33    };
34    let max_count = adapter_limits.max_compute_workgroups_per_dimension;
35    let hardware_props = HardwareProperties {
36        load_width: 128,
37        // On Apple Silicon, the plane size is 32,
38        // though the minimum and maximum differ.
39        // https://github.com/gpuweb/gpuweb/issues/3950
40        #[cfg(apple_silicon)]
41        plane_size_min: 32,
42        #[cfg(not(apple_silicon))]
43        plane_size_min: adapter_info.subgroup_min_size,
44        #[cfg(apple_silicon)]
45        plane_size_max: 32,
46        #[cfg(not(apple_silicon))]
47        plane_size_max: adapter_info.subgroup_max_size,
48        // wgpu uses an additional buffer for variable-length buffers,
49        // so we have to use one buffer less on our side to make room for that wgpu internal buffer.
50        // See: https://github.com/gfx-rs/wgpu/blob/a9638c8e3ac09ce4f27ac171f8175671e30365fd/wgpu-hal/src/metal/device.rs#L799
51        max_bindings: limits
52            .max_storage_buffers_per_shader_stage
53            .saturating_sub(1),
54        max_shared_memory_size: limits.max_compute_workgroup_storage_size as usize,
55        max_ruda_count: (max_count, max_count, max_count),
56        max_units_per_ruda: adapter_limits.max_compute_invocations_per_workgroup,
57        max_ruda_dim: (
58            adapter_limits.max_compute_workgroup_size_x,
59            adapter_limits.max_compute_workgroup_size_y,
60            adapter_limits.max_compute_workgroup_size_z,
61        ),
62        num_streaming_multiprocessors: None,
63        num_tensor_cores: None,
64        min_tensor_cores_dim: None,
65        num_cpu_cores: None, // TODO: Check if device is CPU.
66        max_vector_size: 4,
67    };
68
69    let mut compilation_options = Default::default();
70
71    let features = setup.adapter.features();
72
73    let time_measurement = if features.contains(wgpu::Features::TIMESTAMP_QUERY) {
74        TimingMethod::Device
75    } else {
76        TimingMethod::System
77    };
78
79    let mut device_props = DeviceProperties::new(
80        Default::default(),
81        mem_props,
82        hardware_props,
83        time_measurement,
84    );
85
86    #[cfg(not(all(target_os = "macos", feature = "msl")))]
87    {
88        if features.contains(wgpu::Features::SUBGROUP)
89            && setup.adapter.get_info().device_type != wgpu::DeviceType::Cpu
90        {
91            use ruda_core::ir::features::Plane;
92
93            device_props.features.plane.insert(Plane::Ops);
94        }
95    }
96
97    #[cfg(any(feature = "spirv", feature = "msl"))]
98    device_props
99        .features
100        .plane
101        .insert(ruda_core::ir::features::Plane::NonUniformControlFlow);
102
103    backend::register_features(
104        &setup.adapter,
105        &mut device_props,
106        &mut compilation_options,
107        &options.memory_config,
108    );
109
110    let logger = alloc::sync::Arc::new(ServerLogger::default());
111
112    let allocator = ContiguousMemoryLayoutPolicy::new(device_props.memory.alignment as usize);
113    WgpuServer::new(
114        device_props.memory.clone(),
115        options.memory_config,
116        compilation_options,
117        setup.device.clone(),
118        setup.queue,
119        options.tasks_max,
120        setup.backend,
121        time_measurement,
122        ServerUtilities::new(device_props, logger, setup.backend, allocator),
123    )
124}
ruda_driver_wgpu/runtime/device_service.rs

ruda_driver_wgpu/runtime/
device_service.rs
