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

runmat_plot/gpu/
contour.rs
