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