dynamis_gpu/
context.rs

1use crate::{BindingKind, BindingSpec, ComputePipeline};
2use std::collections::HashMap;
3use std::sync::{Arc, Mutex};
4use wgpu::{
5    Adapter, Backends, Device, Instance, InstanceDescriptor, MemoryHints, PowerPreference, Queue,
6};
7
8const BACKEND_PRIORITY: [Backends; 2] = [Backends::DX12.union(Backends::METAL), Backends::VULKAN];
9
10#[derive(PartialEq, Eq, Hash, Clone)]
11struct PipelineKey {
12    label: String,
13    shader: String,
14    entry: String,
15    workgroup_size: u32,
16    bindings: Vec<(u32, BindingKind)>,
17}
18
19impl PipelineKey {
20    fn of(
21        label: &str,
22        shader: &str,
23        entry: &str,
24        groups: &[&[BindingSpec]],
25        workgroup_size: u32,
26    ) -> Self {
27        let bindings = groups
28            .iter()
29            .flat_map(|group| group.iter())
30            .map(|spec| (spec.binding, spec.kind))
31            .collect();
32        Self {
33            label: label.to_owned(),
34            shader: shader.to_owned(),
35            entry: entry.to_owned(),
36            workgroup_size,
37            bindings,
38        }
39    }
40}
41
42pub struct GpuContext {
43    adapter: Adapter,
44    device: Device,
45    queue: Queue,
46    pipelines: Arc<Mutex<HashMap<PipelineKey, Arc<ComputePipeline>>>>,
47}
48
49impl GpuContext {
50    pub async fn new() -> Self {
51        let adapter = request_adapter().await;
52        let (device, queue) = adapter
53            .request_device(&wgpu::DeviceDescriptor {
54                label: Some("dynamis device"),
55                required_features: wgpu::Features::empty(),
56                required_limits: adapter.limits(),
57                memory_hints: MemoryHints::default(),
58                ..Default::default()
59            })
60            .await
61            .expect("failed to create GPU device");
62        Self {
63            adapter,
64            device,
65            queue,
66            pipelines: Arc::new(Mutex::new(HashMap::new())),
67        }
68    }
69
70    pub fn device(&self) -> &Device {
71        &self.device
72    }
73
74    pub fn queue(&self) -> &Queue {
75        &self.queue
76    }
77
78    pub fn adapter_info(&self) -> wgpu::AdapterInfo {
79        self.adapter.get_info()
80    }
81
82    pub fn compute_pipeline(
83        &self,
84        label: &str,
85        shader: &str,
86        entry: &str,
87        groups: &[&[BindingSpec]],
88        workgroup_size: u32,
89    ) -> ComputePipeline {
90        let key = PipelineKey::of(label, shader, entry, groups, workgroup_size);
91        let mut cache = self.pipelines.lock().unwrap();
92        if let Some(pipeline) = cache.get(&key) {
93            return (**pipeline).clone();
94        }
95        let pipeline = Arc::new(ComputePipeline::new(
96            &self.device,
97            label,
98            shader,
99            entry,
100            groups,
101            workgroup_size,
102        ));
103        cache.insert(key, pipeline.clone());
104        (*pipeline).clone()
105    }
106}
107
108impl Clone for GpuContext {
109    fn clone(&self) -> Self {
110        Self {
111            adapter: self.adapter.clone(),
112            device: self.device.clone(),
113            queue: self.queue.clone(),
114            pipelines: self.pipelines.clone(),
115        }
116    }
117}
118
119async fn request_adapter() -> Adapter {
120    for backends in BACKEND_PRIORITY {
121        let instance = Instance::new(InstanceDescriptor {
122            backends,
123            ..InstanceDescriptor::new_without_display_handle()
124        });
125        let options = wgpu::RequestAdapterOptions {
126            power_preference: PowerPreference::HighPerformance,
127            ..Default::default()
128        };
129        if let Ok(adapter) = instance.request_adapter(&options).await {
130            return adapter;
131        }
132    }
133    panic!("no compatible GPU adapter found for backends {BACKEND_PRIORITY:?}");
134}
dynamis_gpu/context.rs

dynamis_gpu/
context.rs
