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 | wgpu::BufferUsages::COPY_SRC,
105 mapped_at_creation: false,
106 }));
107
108 let indirect_args = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
109 label: Some("contourf-gpu-indirect-args"),
110 size: std::mem::size_of::<DrawIndirectArgsRaw>() as u64,
111 usage: wgpu::BufferUsages::STORAGE
112 | wgpu::BufferUsages::INDIRECT
113 | wgpu::BufferUsages::COPY_DST
114 | wgpu::BufferUsages::COPY_SRC,
115 mapped_at_creation: false,
116 }));
117 let init = DrawIndirectArgsRaw {
118 vertex_count: 0,
119 instance_count: 1,
120 first_vertex: 0,
121 first_instance: 0,
122 };
123 queue.write_buffer(&indirect_args, 0, bytemuck::bytes_of(&init));
124
125 let uniforms = ContourFillUniforms {
126 base_z: params.base_z,
127 alpha: params.alpha,
128 x_len: inputs.x_len,
129 y_len: inputs.y_len,
130 color_table_len: inputs.color_table.len() as u32,
131 band_count,
132 cell_count,
133 _pad: 0,
134 };
135
136 let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
137 label: Some("contourf-uniforms"),
138 contents: bytemuck::bytes_of(&uniforms),
139 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
140 });
141
142 let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
143 label: Some("contourf-bind-layout"),
144 entries: &[
145 wgpu::BindGroupLayoutEntry {
146 binding: 0,
147 visibility: wgpu::ShaderStages::COMPUTE,
148 ty: wgpu::BindingType::Buffer {
149 ty: wgpu::BufferBindingType::Storage { read_only: true },
150 has_dynamic_offset: false,
151 min_binding_size: None,
152 },
153 count: None,
154 },
155 wgpu::BindGroupLayoutEntry {
156 binding: 1,
157 visibility: wgpu::ShaderStages::COMPUTE,
158 ty: wgpu::BindingType::Buffer {
159 ty: wgpu::BufferBindingType::Storage { read_only: true },
160 has_dynamic_offset: false,
161 min_binding_size: None,
162 },
163 count: None,
164 },
165 wgpu::BindGroupLayoutEntry {
166 binding: 2,
167 visibility: wgpu::ShaderStages::COMPUTE,
168 ty: wgpu::BindingType::Buffer {
169 ty: wgpu::BufferBindingType::Storage { read_only: true },
170 has_dynamic_offset: false,
171 min_binding_size: None,
172 },
173 count: None,
174 },
175 wgpu::BindGroupLayoutEntry {
176 binding: 3,
177 visibility: wgpu::ShaderStages::COMPUTE,
178 ty: wgpu::BindingType::Buffer {
179 ty: wgpu::BufferBindingType::Storage { read_only: true },
180 has_dynamic_offset: false,
181 min_binding_size: None,
182 },
183 count: None,
184 },
185 wgpu::BindGroupLayoutEntry {
186 binding: 4,
187 visibility: wgpu::ShaderStages::COMPUTE,
188 ty: wgpu::BindingType::Buffer {
189 ty: wgpu::BufferBindingType::Storage { read_only: true },
190 has_dynamic_offset: false,
191 min_binding_size: None,
192 },
193 count: None,
194 },
195 wgpu::BindGroupLayoutEntry {
196 binding: 5,
197 visibility: wgpu::ShaderStages::COMPUTE,
198 ty: wgpu::BindingType::Buffer {
199 ty: wgpu::BufferBindingType::Storage { read_only: false },
200 has_dynamic_offset: false,
201 min_binding_size: None,
202 },
203 count: None,
204 },
205 wgpu::BindGroupLayoutEntry {
206 binding: 6,
207 visibility: wgpu::ShaderStages::COMPUTE,
208 ty: wgpu::BindingType::Buffer {
209 ty: wgpu::BufferBindingType::Uniform,
210 has_dynamic_offset: false,
211 min_binding_size: None,
212 },
213 count: None,
214 },
215 wgpu::BindGroupLayoutEntry {
216 binding: 7,
217 visibility: wgpu::ShaderStages::COMPUTE,
218 ty: wgpu::BindingType::Buffer {
219 ty: wgpu::BufferBindingType::Storage { read_only: false },
220 has_dynamic_offset: false,
221 min_binding_size: None,
222 },
223 count: None,
224 },
225 ],
226 });
227
228 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
229 label: Some("contourf-pipeline-layout"),
230 bind_group_layouts: &[&bind_group_layout],
231 push_constant_ranges: &[],
232 });
233
234 let pipeline =
235 device.create_compute_pipeline(&crate::wgpu_compat::wgpu_compute_pipeline_descriptor! {
236 label: Some("contourf-pack-pipeline"),
237 layout: Some(&pipeline_layout),
238 module: &shader,
239 entry_point: "main",
240 });
241
242 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
243 label: Some("contourf-bind-group"),
244 layout: &bind_group_layout,
245 entries: &[
246 wgpu::BindGroupEntry {
247 binding: 0,
248 resource: x_buffer.as_ref().as_entire_binding(),
249 },
250 wgpu::BindGroupEntry {
251 binding: 1,
252 resource: y_buffer.as_ref().as_entire_binding(),
253 },
254 wgpu::BindGroupEntry {
255 binding: 2,
256 resource: inputs.z_buffer.as_ref().as_entire_binding(),
257 },
258 wgpu::BindGroupEntry {
259 binding: 3,
260 resource: color_buffer.as_ref().as_entire_binding(),
261 },
262 wgpu::BindGroupEntry {
263 binding: 4,
264 resource: level_buffer.as_ref().as_entire_binding(),
265 },
266 wgpu::BindGroupEntry {
267 binding: 5,
268 resource: output_buffer.as_ref().as_entire_binding(),
269 },
270 wgpu::BindGroupEntry {
271 binding: 6,
272 resource: uniform_buffer.as_entire_binding(),
273 },
274 wgpu::BindGroupEntry {
275 binding: 7,
276 resource: indirect_args.as_ref().as_entire_binding(),
277 },
278 ],
279 });
280
281 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
282 label: Some("contourf-encoder"),
283 });
284 {
285 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
286 label: Some("contourf-pass"),
287 timestamp_writes: None,
288 });
289 pass.set_pipeline(&pipeline);
290 pass.set_bind_group(0, &bind_group, &[]);
291 let workgroups = total_invocations.div_ceil(workgroup_size).max(1);
292 pass.dispatch_workgroups(workgroups, 1, 1);
293 }
294 queue.submit(Some(encoder.finish()));
295
296 Ok(GpuVertexBuffer::with_indirect(
297 output_buffer,
298 vertex_count as usize,
299 indirect_args,
300 ))
301}
302
303fn compile_shader(
304 device: &Arc<wgpu::Device>,
305 workgroup_size: u32,
306 scalar: ScalarType,
307) -> wgpu::ShaderModule {
308 let template = match scalar {
309 ScalarType::F32 => shaders::contour_fill::F32,
310 ScalarType::F64 => shaders::contour_fill::F64,
311 };
312 let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
313 device.create_shader_module(wgpu::ShaderModuleDescriptor {
314 label: Some("contourf-pack-shader"),
315 source: wgpu::ShaderSource::Wgsl(source.into()),
316 })
317}