Skip to main content

pebble/graphics/pipeline/
compute.rs

1use crate::{
2    assets::{handle::Handle, storage::Assets, upload::{Asset, AssetSource}},
3    ecs::resources::Read,
4    graphics::{
5        pipeline::{
6            binding::{BindGroupLayout, BindGroupLayoutBuilder, BindGroupTarget, BindingEntry},
7            layout::{assemble_group_layouts, find_own_entries, GlobalLayoutPool, GroupEntry, PipelineKind},
8        },
9        render::Backend,
10        types::flags::ShaderStages,
11    },
12};
13
14/// A compiled GPU compute pipeline, wrapping `wgpu::ComputePipeline`.
15pub struct ComputePipeline(wgpu::ComputePipeline);
16
17impl ComputePipeline {
18    pub(crate) fn raw(&self) -> &wgpu::ComputePipeline {
19        &self.0
20    }
21}
22
23/// A compute pipeline asset — WGSL shader source plus its bind group
24/// layout. Dispatch it via [`ComputeInstance`](super::instance::ComputeInstance)
25/// and [`Backend::dispatch_compute`](crate::graphics::render::Backend::dispatch_compute).
26pub struct Compute {
27    label: Option<&'static str>,
28    shader_source: &'static str,
29    entry_point: Option<&'static str>,
30    groups: Vec<GroupEntry>,
31}
32
33impl Default for Compute {
34    fn default() -> Self {
35        Self { label: None, shader_source: "", entry_point: Some("cs_main"), groups: Vec::new() }
36    }
37}
38
39impl Compute {
40    pub fn new(shader_source: &'static str) -> Self {
41        Self { shader_source, ..Self::default() }
42    }
43
44    pub fn with_label(mut self, label: &'static str) -> Self {
45        self.label = Some(label);
46        self
47    }
48
49    pub fn with_entry_point(mut self, entry: &'static str) -> Self {
50        self.entry_point = Some(entry);
51        self
52    }
53
54    pub fn with_entries(mut self, groups: Vec<GroupEntry>) -> Self {
55        self.groups = groups;
56        self
57    }
58
59    fn validate(&self) {
60        if self.groups.is_empty() {
61            tracing::warn!(
62                "Compute{}: no bind groups at all — this pass can't read or write \
63                 anything; consider calling .with_entries(...)",
64                self.label.map(|l| format!(" '{l}'")).unwrap_or_default(),
65            );
66        }
67    }
68
69    pub fn build_asset(self, name: &str, assets: &mut Assets<Compute>) -> Handle<Compute> {
70        self.validate();
71        assets.insert(name, self)
72    }
73}
74
75/// Compiles a [`Compute`] into a raw pipeline + bind group layout. Used
76/// internally by the asset upload path; exposed for callers assembling
77/// pipelines outside the usual [`Assets`] flow.
78pub fn build_compute(backend: &Backend, desc: &Compute, pool: &GlobalLayoutPool) -> Option<(ComputePipeline, BindGroupLayout)> {
79    let own_entries = find_own_entries(desc.label, PipelineKind::Compute, &desc.groups);
80    for entry in own_entries {
81        if entry.kind.visibility() != ShaderStages::COMPUTE {
82            panic!(
83                "compute pass{}: entry '{}' is not visible to exactly the compute stage — \
84                 compute bind group entries must be visible to exactly COMPUTE",
85                desc.label.map(|l| format!(" '{l}'")).unwrap_or_default(),
86                entry.name,
87            );
88        }
89    }
90
91    let layout = BindGroupLayoutBuilder::new()
92        .with_label(desc.label)
93        .with_entries(own_entries.iter().cloned())
94        .build(backend);
95
96    let device = &backend.device;
97    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
98        label: desc.label,
99        source: wgpu::ShaderSource::Wgsl(desc.shader_source.into()),
100    });
101
102    let bind_group_layouts = assemble_group_layouts(
103        desc.label,
104        &desc.groups,
105        &layout,
106        pool,
107        device.limits().max_bind_groups,
108    )?;
109
110    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
111        label: desc.label,
112        bind_group_layouts: &bind_group_layouts,
113        immediate_size: 0,
114    });
115
116    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
117        label: desc.label,
118        layout: Some(&pipeline_layout),
119        module: &module,
120        entry_point: desc.entry_point,
121        compilation_options: Default::default(),
122        cache: None,
123    });
124
125    Some((ComputePipeline(pipeline), layout))
126}
127
128/// The GPU-resident pipeline an uploaded [`Compute`] produces.
129pub struct GPUCompute {
130    pub pipeline: ComputePipeline,
131    layout: BindGroupLayout,
132    entries: Vec<BindingEntry>,
133}
134
135impl BindGroupTarget for GPUCompute {
136    fn bind_group_layout(&self) -> &BindGroupLayout {
137        &self.layout
138    }
139    fn binding_entries(&self) -> &[BindingEntry] {
140        &self.entries
141    }
142}
143
144impl AssetSource for Compute {
145    type Processed = GPUCompute;
146}
147
148impl Asset<Backend> for Compute {
149    type Deps<'a> = Read<'a, GlobalLayoutPool>;
150
151    fn upload<'a>(&self, backend: &Backend, pool: &Read<'a, GlobalLayoutPool>) -> Option<GPUCompute> {
152        let (pipeline, layout) = build_compute(backend, self, pool)?;
153        let entries = find_own_entries(self.label, PipelineKind::Compute, &self.groups).to_vec();
154
155        Some(GPUCompute { pipeline, layout, entries })
156    }
157}