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

runmat_plot/gpu/
contour.rs
