runmat_plot/gpu/
contour.rs

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

runmat_plot/gpu/
contour.rs
