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 =
233 device.create_compute_pipeline(&crate::wgpu_compat::wgpu_compute_pipeline_descriptor! {
234 label: Some("contourf-pack-pipeline"),
235 layout: Some(&pipeline_layout),
236 module: &shader,
237 entry_point: "main",
238 });
239
240 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
241 label: Some("contourf-bind-group"),
242 layout: &bind_group_layout,
243 entries: &[
244 wgpu::BindGroupEntry {
245 binding: 0,
246 resource: x_buffer.as_ref().as_entire_binding(),
247 },
248 wgpu::BindGroupEntry {
249 binding: 1,
250 resource: y_buffer.as_ref().as_entire_binding(),
251 },
252 wgpu::BindGroupEntry {
253 binding: 2,
254 resource: inputs.z_buffer.as_ref().as_entire_binding(),
255 },
256 wgpu::BindGroupEntry {
257 binding: 3,
258 resource: color_buffer.as_ref().as_entire_binding(),
259 },
260 wgpu::BindGroupEntry {
261 binding: 4,
262 resource: level_buffer.as_ref().as_entire_binding(),
263 },
264 wgpu::BindGroupEntry {
265 binding: 5,
266 resource: output_buffer.as_ref().as_entire_binding(),
267 },
268 wgpu::BindGroupEntry {
269 binding: 6,
270 resource: uniform_buffer.as_entire_binding(),
271 },
272 wgpu::BindGroupEntry {
273 binding: 7,
274 resource: indirect_args.as_ref().as_entire_binding(),
275 },
276 ],
277 });
278
279 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
280 label: Some("contourf-encoder"),
281 });
282 {
283 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
284 label: Some("contourf-pass"),
285 timestamp_writes: None,
286 });
287 pass.set_pipeline(&pipeline);
288 pass.set_bind_group(0, &bind_group, &[]);
289 let workgroups = total_invocations.div_ceil(workgroup_size).max(1);
290 pass.dispatch_workgroups(workgroups, 1, 1);
291 }
292 queue.submit(Some(encoder.finish()));
293
294 Ok(GpuVertexBuffer::with_indirect(
295 output_buffer,
296 vertex_count as usize,
297 indirect_args,
298 ))
299}
300
301fn compile_shader(
302 device: &Arc<wgpu::Device>,
303 workgroup_size: u32,
304 scalar: ScalarType,
305) -> wgpu::ShaderModule {
306 let template = match scalar {
307 ScalarType::F32 => shaders::contour_fill::F32,
308 ScalarType::F64 => shaders::contour_fill::F64,
309 };
310 let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
311 device.create_shader_module(wgpu::ShaderModuleDescriptor {
312 label: Some("contourf-pack-shader"),
313 source: wgpu::ShaderSource::Wgsl(source.into()),
314 })
315}