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 = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
216        label: Some("contour-pack-pipeline"),
217        layout: Some(&pipeline_layout),
218        module: &shader,
219        entry_point: "main",
220    });
221
222    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
223        label: Some("contour-pack-bind-group"),
224        layout: &bind_group_layout,
225        entries: &[
226            wgpu::BindGroupEntry {
227                binding: 0,
228                resource: x_buffer.as_ref().as_entire_binding(),
229            },
230            wgpu::BindGroupEntry {
231                binding: 1,
232                resource: y_buffer.as_ref().as_entire_binding(),
233            },
234            wgpu::BindGroupEntry {
235                binding: 2,
236                resource: inputs.z_buffer.as_ref().as_entire_binding(),
237            },
238            wgpu::BindGroupEntry {
239                binding: 3,
240                resource: color_buffer.as_ref().as_entire_binding(),
241            },
242            wgpu::BindGroupEntry {
243                binding: 4,
244                resource: output_buffer.as_ref().as_entire_binding(),
245            },
246            wgpu::BindGroupEntry {
247                binding: 5,
248                resource: uniform_buffer.as_entire_binding(),
249            },
250            wgpu::BindGroupEntry {
251                binding: 6,
252                resource: level_buffer.as_ref().as_entire_binding(),
253            },
254        ],
255    });
256
257    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
258        label: Some("contour-pack-encoder"),
259    });
260    {
261        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
262            label: Some("contour-pack-pass"),
263            timestamp_writes: None,
264        });
265        pass.set_pipeline(&pipeline);
266        pass.set_bind_group(0, &bind_group, &[]);
267        let workgroups = total_invocations.div_ceil(workgroup_size);
268        pass.dispatch_workgroups(workgroups, 1, 1);
269    }
270    queue.submit(Some(encoder.finish()));
271
272    Ok(GpuVertexBuffer::new(output_buffer, vertex_count as usize))
273}
274
275fn compile_shader(
276    device: &Arc<wgpu::Device>,
277    workgroup_size: u32,
278    scalar: ScalarType,
279) -> wgpu::ShaderModule {
280    let template = match scalar {
281        ScalarType::F32 => shaders::contour::F32,
282        ScalarType::F64 => shaders::contour::F64,
283    };
284    let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
285    device.create_shader_module(wgpu::ShaderModuleDescriptor {
286        label: Some("contour-pack-shader"),
287        source: wgpu::ShaderSource::Wgsl(source.into()),
288    })
289}
runmat_plot/gpu/contour.rs

runmat_plot/gpu/
contour.rs
