Skip to main content

dynamis_gpu/
library.rs

1use crate::pipeline::{ComputeLayout, ComputeProgram, PipelineHandle};
2use std::collections::{HashMap, VecDeque};
3use std::sync::Arc;
4use std::time::Duration;
5use wgpu::Device;
6
7#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8pub enum WarmupBudget {
9    All,
10    Within(Duration),
11}
12
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14pub struct WarmupProgress {
15    pub ready: usize,
16    pub total: usize,
17}
18
19impl WarmupProgress {
20    pub fn complete(&self) -> bool {
21        self.ready == self.total
22    }
23}
24
25pub(crate) struct PipelineLibrary {
26    entries: HashMap<Arc<ComputeProgram>, PipelineHandle>,
27    pending: VecDeque<PipelineHandle>,
28}
29
30impl PipelineLibrary {
31    pub(crate) fn new() -> Self {
32        Self {
33            entries: HashMap::new(),
34            pending: VecDeque::new(),
35        }
36    }
37
38    pub(crate) fn declare(&mut self, device: &Device, program: ComputeProgram) -> PipelineHandle {
39        let program = Arc::new(program);
40        if let Some(handle) = self.entries.get(&program) {
41            return handle.clone();
42        }
43        let handle = PipelineHandle::new(program.clone(), ComputeLayout::new(device, &program));
44        self.entries.insert(program, handle.clone());
45        self.pending.push_back(handle.clone());
46        handle
47    }
48
49    pub(crate) fn drain_pending(&mut self) -> Vec<PipelineHandle> {
50        self.pending.drain(..).collect()
51    }
52
53    pub(crate) fn take_pending(&mut self) -> Option<PipelineHandle> {
54        self.pending.pop_front()
55    }
56
57    pub(crate) fn progress(&self) -> WarmupProgress {
58        WarmupProgress {
59            ready: self
60                .entries
61                .values()
62                .filter(|handle| handle.is_warmed())
63                .count(),
64            total: self.entries.len(),
65        }
66    }
67}