1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83

use std::sync::Arc;

use crate::{
    bind_group::BindGroup, context::Context, pipeline_layout::PipelineLayout, shader::EntryPoint,
};

/// A compute pipeline
///
/// Loosely equivalent to [wgpu::ComputePipeline]
#[derive(Clone, Debug)]
pub struct ComputePipeline {
    entry_point: EntryPoint,
    label: Option<String>,
}

#[derive(Clone, Hash, PartialEq, Eq)]
pub(crate) struct ComputePipelineCacheKey {
    layout: PipelineLayout,
    entry_point: EntryPoint,
}

impl ComputePipeline {
    pub(crate) fn get_or_build(
        &self,
        context: &Context,
        bind_groups: &[BindGroup],
    ) -> Arc<wgpu::ComputePipeline> {
        let layout = PipelineLayout {
            bind_group_layouts: bind_groups.iter().map(|b| b.build_layout()).collect(),
        };

        let key = ComputePipelineCacheKey {
            layout: layout.clone(),
            entry_point: self.entry_point.clone(),
        };

        let mut pipeline_cache = context.ctx.caches.compute_pipeline_cache.borrow_mut();

        pipeline_cache
            .get_or_insert_with(key, || {
                let layout = layout.get_or_build(context);

                Arc::new(context.device().create_compute_pipeline(
                    &wgpu::ComputePipelineDescriptor {
                        layout: Some(&layout),
                        module: &self.entry_point.shader,
                        entry_point: &self.entry_point.entry_point,
                        label: self.label.as_deref(),
                    },
                ))
            })
            .clone()
    }
}

/// Builds a [ComputePipeline]
#[derive(Clone)]
pub struct ComputePipelineBuilder {
    entry_point: EntryPoint,
    label: Option<String>,
}

impl ComputePipelineBuilder {
    pub fn with_entry_point(entry_point: &EntryPoint) -> Self {
        Self {
            entry_point: entry_point.clone(),
            label: None,
        }
    }

    /// Set the optional debug name. This may appear in error messages and GPU profiler traces
    pub fn label(mut self, label: &str) -> Self {
        self.label = Some(label.into());
        self
    }

    pub fn build(self) -> ComputePipeline {
        ComputePipeline {
            entry_point: self.entry_point,
            label: self.label,
        }
    }
}