pebble/wgpu/
compute.rs

1use crate::{
2    app::App,
3    assets::{plugin::AssetPlugin, upload::Asset},
4    ecs::plugin::Plugin,
5    wgpu::backend::WGPUBackend,
6};
7
8#[derive(Copy, Clone, PartialEq, Eq, Hash)]
9pub enum ComputeBindingKind {
10    StorageBufferReadOnly,
11    StorageBufferReadWrite,
12    UniformBuffer,
13    Texture,
14    StorageTexture {
15        format: wgpu::TextureFormat,
16        access: wgpu::StorageTextureAccess,
17    },
18    Sampler,
19}
20
21impl ComputeBindingKind {
22    pub fn layout_entry(&self, binding: u32) -> wgpu::BindGroupLayoutEntry {
23        match self {
24            ComputeBindingKind::StorageBufferReadOnly => wgpu::BindGroupLayoutEntry {
25                binding,
26                visibility: wgpu::ShaderStages::COMPUTE,
27                ty: wgpu::BindingType::Buffer {
28                    ty: wgpu::BufferBindingType::Storage { read_only: true },
29                    has_dynamic_offset: false,
30                    min_binding_size: None,
31                },
32                count: None,
33            },
34            ComputeBindingKind::StorageBufferReadWrite => wgpu::BindGroupLayoutEntry {
35                binding,
36                visibility: wgpu::ShaderStages::COMPUTE,
37                ty: wgpu::BindingType::Buffer {
38                    ty: wgpu::BufferBindingType::Storage { read_only: false },
39                    has_dynamic_offset: false,
40                    min_binding_size: None,
41                },
42                count: None,
43            },
44            ComputeBindingKind::UniformBuffer => wgpu::BindGroupLayoutEntry {
45                binding,
46                visibility: wgpu::ShaderStages::COMPUTE,
47                ty: wgpu::BindingType::Buffer {
48                    ty: wgpu::BufferBindingType::Uniform,
49                    has_dynamic_offset: false,
50                    min_binding_size: None,
51                },
52                count: None,
53            },
54            ComputeBindingKind::Texture => wgpu::BindGroupLayoutEntry {
55                binding,
56                visibility: wgpu::ShaderStages::COMPUTE,
57                ty: wgpu::BindingType::Texture {
58                    sample_type: wgpu::TextureSampleType::Float { filterable: true },
59                    view_dimension: wgpu::TextureViewDimension::D2,
60                    multisampled: false,
61                },
62                count: None,
63            },
64            ComputeBindingKind::StorageTexture { format, access } => wgpu::BindGroupLayoutEntry {
65                binding,
66                visibility: wgpu::ShaderStages::COMPUTE,
67                ty: wgpu::BindingType::StorageTexture {
68                    access: *access,
69                    format: *format,
70                    view_dimension: wgpu::TextureViewDimension::D2,
71                },
72                count: None,
73            },
74            ComputeBindingKind::Sampler => wgpu::BindGroupLayoutEntry {
75                binding,
76                visibility: wgpu::ShaderStages::COMPUTE,
77                ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
78                count: None,
79            },
80        }
81    }
82}
83
84#[derive(Clone)]
85pub struct ComputeBindingEntry {
86    pub name: &'static str,
87    pub kind: ComputeBindingKind,
88}
89
90pub struct ComputeDescriptor<'a> {
91    pub label: Option<&'a str>,
92    pub shader_source: &'a str,
93    pub entry_point: Option<&'a str>,
94    pub entries: Vec<ComputeBindingEntry>,
95    pub extra_layouts: Vec<wgpu::BindGroupLayout>,
96}
97
98impl<'a> Default for ComputeDescriptor<'a> {
99    fn default() -> Self {
100        Self {
101            label: None,
102            shader_source: "",
103            entry_point: Some("cs_main"),
104            entries: Vec::new(),
105            extra_layouts: Vec::new(),
106        }
107    }
108}
109
110pub fn build_compute(
111    device: &wgpu::Device,
112    desc: &ComputeDescriptor,
113) -> (wgpu::ComputePipeline, wgpu::BindGroupLayout) {
114    let layout_entries: Vec<_> = desc
115        .entries
116        .iter()
117        .enumerate()
118        .map(|(i, e)| e.kind.layout_entry(i as u32))
119        .collect();
120
121    let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
122        label: desc.label,
123        entries: &layout_entries,
124    });
125
126    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
127        label: desc.label,
128        source: wgpu::ShaderSource::Wgsl(desc.shader_source.into()),
129    });
130
131    let mut bind_group_layouts: Vec<&wgpu::BindGroupLayout> = desc.extra_layouts.iter().collect();
132    bind_group_layouts.push(&layout);
133    let bind_group_layouts: Vec<Option<&wgpu::BindGroupLayout>> =
134        bind_group_layouts.into_iter().map(Some).collect();
135
136    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
137        label: desc.label,
138        bind_group_layouts: &bind_group_layouts,
139        immediate_size: 0,
140    });
141
142    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
143        label: desc.label,
144        layout: Some(&pipeline_layout),
145        module: &module,
146        entry_point: desc.entry_point,
147        compilation_options: Default::default(),
148        cache: None,
149    });
150
151    (pipeline, layout)
152}
153
154pub struct GPUCompute {
155    pub pipeline: wgpu::ComputePipeline,
156    pub layout: wgpu::BindGroupLayout,
157    pub entries: Vec<ComputeBindingEntry>,
158}
159
160impl Asset<WGPUBackend> for GPUCompute {
161    type Source = ComputeDescriptor<'static>;
162    type Deps<'a> = ();
163
164    fn upload<'a>(source: &ComputeDescriptor, backend: &WGPUBackend, _deps: &()) -> Option<Self> {
165        let (pipeline, layout) = build_compute(&backend.device, source);
166
167        Some(Self {
168            pipeline,
169            layout,
170            entries: source.entries.to_vec(),
171        })
172    }
173}
174
175pub struct ComputePlugin;
176impl ComputePlugin {
177    pub fn new() -> Self {
178        Self
179    }
180}
181
182impl Plugin for ComputePlugin {
183    fn build(&self, app: &mut App) {
184        app.add_plugin(AssetPlugin::<super::backend::WGPUBackend, GPUCompute>::new());
185    }
186}
pebble/wgpu/compute.rs

pebble/wgpu/
compute.rs
