1use crate::core::renderer::Vertex;
2use crate::core::scene::{DrawIndirectArgsRaw, GpuVertexBuffer};
3use crate::gpu::axis::{axis_storage_buffer, AxisData};
4use crate::gpu::shaders;
5use crate::gpu::{tuning, ScalarType};
6use std::sync::Arc;
7use wgpu::util::DeviceExt;
8
9pub struct ContourFillGpuInputs<'a> {
10 pub x_axis: AxisData<'a>,
11 pub y_axis: AxisData<'a>,
12 pub z_buffer: Arc<wgpu::Buffer>,
13 pub color_table: &'a [[f32; 4]],
14 pub level_values: &'a [f32],
15 pub x_len: u32,
16 pub y_len: u32,
17 pub scalar: ScalarType,
18}
19
20pub struct ContourFillGpuParams {
21 pub base_z: f32,
22 pub alpha: f32,
23}
24
25#[repr(C)]
26#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
27struct ContourFillUniforms {
28 base_z: f32,
29 alpha: f32,
30 x_len: u32,
31 y_len: u32,
32 color_table_len: u32,
33 band_count: u32,
34 cell_count: u32,
35 _pad: u32,
36}
37
38const MAX_VERTICES_PER_INVOCATION: u32 = 36;
39
40pub fn pack_contour_fill_vertices(
41 device: &Arc<wgpu::Device>,
42 queue: &Arc<wgpu::Queue>,
43 inputs: &ContourFillGpuInputs<'_>,
44 params: &ContourFillGpuParams,
45) -> Result<GpuVertexBuffer, String> {
46 if inputs.x_len < 2 || inputs.y_len < 2 {
47 return Err("contourf: axis vectors must contain at least two elements".to_string());
48 }
49 if inputs.level_values.len() < 2 {
50 return Err("contourf: level vector must contain at least two entries".to_string());
51 }
52 if inputs.color_table.is_empty() {
53 return Err("contourf: color table must contain at least one entry".to_string());
54 }
55
56 let cells_x = inputs.x_len - 1;
57 let cells_y = inputs.y_len - 1;
58 let cell_count = cells_x
59 .checked_mul(cells_y)
60 .ok_or_else(|| "contourf: grid dimensions overflowed cell count".to_string())?;
61 if cell_count == 0 {
62 return Err("contourf: no cells available for fill generation".to_string());
63 }
64 let band_count = (inputs.level_values.len() as u32).saturating_sub(1);
65 if band_count == 0 {
66 return Err("contourf: at least one fill band is required".to_string());
67 }
68 let total_invocations = cell_count
69 .checked_mul(band_count)
70 .ok_or_else(|| "contourf: invocation count overflowed".to_string())?;
71 let vertex_count = total_invocations
72 .checked_mul(MAX_VERTICES_PER_INVOCATION)
73 .ok_or_else(|| "contourf: vertex count overflowed".to_string())?;
74
75 let workgroup_size = tuning::effective_workgroup_size();
76 let shader = compile_shader(device, workgroup_size, inputs.scalar);
77
78 let x_buffer = axis_storage_buffer(device, "contourf-x-axis", &inputs.x_axis, inputs.scalar)?;
79 let y_buffer = axis_storage_buffer(device, "contourf-y-axis", &inputs.y_axis, inputs.scalar)?;
80
81 let color_buffer = Arc::new(
82 device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
83 label: Some("contourf-color-table"),
84 contents: bytemuck::cast_slice(inputs.color_table),
85 usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
86 }),
87 );
88
89 let level_buffer = Arc::new(
90 device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
91 label: Some("contourf-level-values"),
92 contents: bytemuck::cast_slice(inputs.level_values),
93 usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
94 }),
95 );
96
97 let output_size = vertex_count as u64 * std::mem::size_of::<Vertex>() as u64;
98 let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
99 label: Some("contourf-gpu-vertices"),
100 size: output_size,
101 usage: wgpu::BufferUsages::STORAGE
102 | wgpu::BufferUsages::VERTEX
103 | wgpu::BufferUsages::COPY_DST,
104 mapped_at_creation: false,
105 }));
106
107 let indirect_args = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
108 label: Some("contourf-gpu-indirect-args"),
109 size: std::mem::size_of::<DrawIndirectArgsRaw>() as u64,
110 usage: wgpu::BufferUsages::STORAGE
111 | wgpu::BufferUsages::INDIRECT
112 | wgpu::BufferUsages::COPY_DST,
113 mapped_at_creation: false,
114 }));
115 let init = DrawIndirectArgsRaw {
116 vertex_count: 0,
117 instance_count: 1,
118 first_vertex: 0,
119 first_instance: 0,
120 };
121 queue.write_buffer(&indirect_args, 0, bytemuck::bytes_of(&init));
122
123 let uniforms = ContourFillUniforms {
124 base_z: params.base_z,
125 alpha: params.alpha,
126 x_len: inputs.x_len,
127 y_len: inputs.y_len,
128 color_table_len: inputs.color_table.len() as u32,
129 band_count,
130 cell_count,
131 _pad: 0,
132 };
133
134 let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
135 label: Some("contourf-uniforms"),
136 contents: bytemuck::bytes_of(&uniforms),
137 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
138 });
139
140 let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
141 label: Some("contourf-bind-layout"),
142 entries: &[
143 wgpu::BindGroupLayoutEntry {
144 binding: 0,
145 visibility: wgpu::ShaderStages::COMPUTE,
146 ty: wgpu::BindingType::Buffer {
147 ty: wgpu::BufferBindingType::Storage { read_only: true },
148 has_dynamic_offset: false,
149 min_binding_size: None,
150 },
151 count: None,
152 },
153 wgpu::BindGroupLayoutEntry {
154 binding: 1,
155 visibility: wgpu::ShaderStages::COMPUTE,
156 ty: wgpu::BindingType::Buffer {
157 ty: wgpu::BufferBindingType::Storage { read_only: true },
158 has_dynamic_offset: false,
159 min_binding_size: None,
160 },
161 count: None,
162 },
163 wgpu::BindGroupLayoutEntry {
164 binding: 2,
165 visibility: wgpu::ShaderStages::COMPUTE,
166 ty: wgpu::BindingType::Buffer {
167 ty: wgpu::BufferBindingType::Storage { read_only: true },
168 has_dynamic_offset: false,
169 min_binding_size: None,
170 },
171 count: None,
172 },
173 wgpu::BindGroupLayoutEntry {
174 binding: 3,
175 visibility: wgpu::ShaderStages::COMPUTE,
176 ty: wgpu::BindingType::Buffer {
177 ty: wgpu::BufferBindingType::Storage { read_only: true },
178 has_dynamic_offset: false,
179 min_binding_size: None,
180 },
181 count: None,
182 },
183 wgpu::BindGroupLayoutEntry {
184 binding: 4,
185 visibility: wgpu::ShaderStages::COMPUTE,
186 ty: wgpu::BindingType::Buffer {
187 ty: wgpu::BufferBindingType::Storage { read_only: true },
188 has_dynamic_offset: false,
189 min_binding_size: None,
190 },
191 count: None,
192 },
193 wgpu::BindGroupLayoutEntry {
194 binding: 5,
195 visibility: wgpu::ShaderStages::COMPUTE,
196 ty: wgpu::BindingType::Buffer {
197 ty: wgpu::BufferBindingType::Storage { read_only: false },
198 has_dynamic_offset: false,
199 min_binding_size: None,
200 },
201 count: None,
202 },
203 wgpu::BindGroupLayoutEntry {
204 binding: 6,
205 visibility: wgpu::ShaderStages::COMPUTE,
206 ty: wgpu::BindingType::Buffer {
207 ty: wgpu::BufferBindingType::Uniform,
208 has_dynamic_offset: false,
209 min_binding_size: None,
210 },
211 count: None,
212 },
213 wgpu::BindGroupLayoutEntry {
214 binding: 7,
215 visibility: wgpu::ShaderStages::COMPUTE,
216 ty: wgpu::BindingType::Buffer {
217 ty: wgpu::BufferBindingType::Storage { read_only: false },
218 has_dynamic_offset: false,
219 min_binding_size: None,
220 },
221 count: None,
222 },
223 ],
224 });
225
226 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
227 label: Some("contourf-pipeline-layout"),
228 bind_group_layouts: &[&bind_group_layout],
229 push_constant_ranges: &[],
230 });
231
232 let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
233 label: Some("contourf-pack-pipeline"),
234 layout: Some(&pipeline_layout),
235 module: &shader,
236 entry_point: "main",
237 });
238
239 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
240 label: Some("contourf-bind-group"),
241 layout: &bind_group_layout,
242 entries: &[
243 wgpu::BindGroupEntry {
244 binding: 0,
245 resource: x_buffer.as_ref().as_entire_binding(),
246 },
247 wgpu::BindGroupEntry {
248 binding: 1,
249 resource: y_buffer.as_ref().as_entire_binding(),
250 },
251 wgpu::BindGroupEntry {
252 binding: 2,
253 resource: inputs.z_buffer.as_ref().as_entire_binding(),
254 },
255 wgpu::BindGroupEntry {
256 binding: 3,
257 resource: color_buffer.as_ref().as_entire_binding(),
258 },
259 wgpu::BindGroupEntry {
260 binding: 4,
261 resource: level_buffer.as_ref().as_entire_binding(),
262 },
263 wgpu::BindGroupEntry {
264 binding: 5,
265 resource: output_buffer.as_ref().as_entire_binding(),
266 },
267 wgpu::BindGroupEntry {
268 binding: 6,
269 resource: uniform_buffer.as_entire_binding(),
270 },
271 wgpu::BindGroupEntry {
272 binding: 7,
273 resource: indirect_args.as_ref().as_entire_binding(),
274 },
275 ],
276 });
277
278 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
279 label: Some("contourf-encoder"),
280 });
281 {
282 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
283 label: Some("contourf-pass"),
284 timestamp_writes: None,
285 });
286 pass.set_pipeline(&pipeline);
287 pass.set_bind_group(0, &bind_group, &[]);
288 let workgroups = total_invocations.div_ceil(workgroup_size).max(1);
289 pass.dispatch_workgroups(workgroups, 1, 1);
290 }
291 queue.submit(Some(encoder.finish()));
292
293 Ok(GpuVertexBuffer::with_indirect(
294 output_buffer,
295 vertex_count as usize,
296 indirect_args,
297 ))
298}
299
300fn compile_shader(
301 device: &Arc<wgpu::Device>,
302 workgroup_size: u32,
303 scalar: ScalarType,
304) -> wgpu::ShaderModule {
305 let template = match scalar {
306 ScalarType::F32 => shaders::contour_fill::F32,
307 ScalarType::F64 => shaders::contour_fill::F64,
308 };
309 let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
310 device.create_shader_module(wgpu::ShaderModuleDescriptor {
311 label: Some("contourf-pack-shader"),
312 source: wgpu::ShaderSource::Wgsl(source.into()),
313 })
314}