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}