runmat_plot/gpu/
surface.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/// Axis data source used by the GPU surface vertex packer.
10pub type SurfaceAxis<'a> = AxisData<'a>;
11
12/// Inputs required to pack surface vertices directly on the GPU.
13pub struct SurfaceGpuInputs<'a> {
14    pub x_axis: SurfaceAxis<'a>,
15    pub y_axis: SurfaceAxis<'a>,
16    pub z_buffer: Arc<wgpu::Buffer>,
17    pub color_table: &'a [[f32; 4]],
18    pub x_len: u32,
19    pub y_len: u32,
20    pub scalar: ScalarType,
21}
22
23/// Parameters describing how the GPU vertices should be generated.
24pub struct SurfaceGpuParams {
25    pub min_z: f32,
26    pub max_z: f32,
27    pub alpha: f32,
28    pub flatten_z: bool,
29    pub x_stride: u32,
30    pub y_stride: u32,
31    pub lod_x_len: u32,
32    pub lod_y_len: u32,
33}
34
35#[repr(C)]
36#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
37struct SurfaceUniforms {
38    min_z: f32,
39    max_z: f32,
40    alpha: f32,
41    flatten: u32,
42    x_len: u32,
43    y_len: u32,
44    lod_x_len: u32,
45    lod_y_len: u32,
46    x_stride: u32,
47    y_stride: u32,
48    color_table_len: u32,
49    _pad: u32,
50}
51
52/// Builds a GPU-resident vertex buffer for surface plots directly from provider-owned Z data.
53pub fn pack_surface_vertices(
54    device: &Arc<wgpu::Device>,
55    queue: &Arc<wgpu::Queue>,
56    inputs: &SurfaceGpuInputs<'_>,
57    params: &SurfaceGpuParams,
58) -> Result<GpuVertexBuffer, String> {
59    if inputs.x_len < 2 || inputs.y_len < 2 {
60        return Err("surf: axis vectors must contain at least two elements".to_string());
61    }
62
63    let workgroup_size = tuning::effective_workgroup_size();
64    let shader = compile_shader(device, workgroup_size, inputs.scalar);
65
66    let x_buffer = axis_storage_buffer(device, "surface-x-axis", &inputs.x_axis, inputs.scalar)?;
67    let y_buffer = axis_storage_buffer(device, "surface-y-axis", &inputs.y_axis, inputs.scalar)?;
68
69    let color_buffer = Arc::new(
70        device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
71            label: Some("surface-color-table"),
72            contents: bytemuck::cast_slice(inputs.color_table),
73            usage: wgpu::BufferUsages::STORAGE
74                | wgpu::BufferUsages::COPY_DST
75                | wgpu::BufferUsages::COPY_SRC,
76        }),
77    );
78
79    let lod_x_len = params.lod_x_len.max(1);
80    let lod_y_len = params.lod_y_len.max(1);
81    let vertex_count = lod_x_len
82        .checked_mul(lod_y_len)
83        .ok_or_else(|| "surf: grid dimensions overflowed vertex count".to_string())?;
84    let output_size = vertex_count as u64 * std::mem::size_of::<Vertex>() as u64;
85    let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
86        label: Some("surface-gpu-vertices"),
87        size: output_size,
88        usage: wgpu::BufferUsages::STORAGE
89            | wgpu::BufferUsages::VERTEX
90            | wgpu::BufferUsages::COPY_DST
91            | wgpu::BufferUsages::COPY_SRC,
92        mapped_at_creation: false,
93    }));
94
95    let min_z = if params.min_z.is_finite() {
96        params.min_z
97    } else {
98        tracing::warn!(
99            target: "runmat_plot::surface_gpu",
100            min_z = params.min_z,
101            "non-finite min_z received; sanitizing to 0.0"
102        );
103        0.0
104    };
105    let mut max_z = if params.max_z.is_finite() {
106        params.max_z
107    } else {
108        tracing::warn!(
109            target: "runmat_plot::surface_gpu",
110            max_z = params.max_z,
111            "non-finite max_z received; sanitizing to min_z + 1.0"
112        );
113        min_z + 1.0
114    };
115    if max_z <= min_z {
116        tracing::warn!(
117            target: "runmat_plot::surface_gpu",
118            min_z,
119            max_z,
120            "invalid z range received; forcing epsilon span"
121        );
122        max_z = min_z + 1e-6;
123    }
124
125    let uniforms = SurfaceUniforms {
126        min_z,
127        max_z,
128        alpha: params.alpha,
129        flatten: if params.flatten_z { 1 } else { 0 },
130        x_len: inputs.x_len,
131        y_len: inputs.y_len,
132        lod_x_len,
133        lod_y_len,
134        x_stride: params.x_stride.max(1),
135        y_stride: params.y_stride.max(1),
136        color_table_len: inputs.color_table.len() as u32,
137        _pad: 0,
138    };
139    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
140        label: Some("surface-pack-uniforms"),
141        contents: bytemuck::bytes_of(&uniforms),
142        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
143    });
144
145    let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
146        label: Some("surface-pack-bind-layout"),
147        entries: &[
148            wgpu::BindGroupLayoutEntry {
149                binding: 0,
150                visibility: wgpu::ShaderStages::COMPUTE,
151                ty: wgpu::BindingType::Buffer {
152                    ty: wgpu::BufferBindingType::Storage { read_only: true },
153                    has_dynamic_offset: false,
154                    min_binding_size: None,
155                },
156                count: None,
157            },
158            wgpu::BindGroupLayoutEntry {
159                binding: 1,
160                visibility: wgpu::ShaderStages::COMPUTE,
161                ty: wgpu::BindingType::Buffer {
162                    ty: wgpu::BufferBindingType::Storage { read_only: true },
163                    has_dynamic_offset: false,
164                    min_binding_size: None,
165                },
166                count: None,
167            },
168            wgpu::BindGroupLayoutEntry {
169                binding: 2,
170                visibility: wgpu::ShaderStages::COMPUTE,
171                ty: wgpu::BindingType::Buffer {
172                    ty: wgpu::BufferBindingType::Storage { read_only: true },
173                    has_dynamic_offset: false,
174                    min_binding_size: None,
175                },
176                count: None,
177            },
178            wgpu::BindGroupLayoutEntry {
179                binding: 3,
180                visibility: wgpu::ShaderStages::COMPUTE,
181                ty: wgpu::BindingType::Buffer {
182                    ty: wgpu::BufferBindingType::Storage { read_only: true },
183                    has_dynamic_offset: false,
184                    min_binding_size: None,
185                },
186                count: None,
187            },
188            wgpu::BindGroupLayoutEntry {
189                binding: 4,
190                visibility: wgpu::ShaderStages::COMPUTE,
191                ty: wgpu::BindingType::Buffer {
192                    ty: wgpu::BufferBindingType::Storage { read_only: false },
193                    has_dynamic_offset: false,
194                    min_binding_size: None,
195                },
196                count: None,
197            },
198            wgpu::BindGroupLayoutEntry {
199                binding: 5,
200                visibility: wgpu::ShaderStages::COMPUTE,
201                ty: wgpu::BindingType::Buffer {
202                    ty: wgpu::BufferBindingType::Uniform,
203                    has_dynamic_offset: false,
204                    min_binding_size: None,
205                },
206                count: None,
207            },
208        ],
209    });
210
211    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
212        label: Some("surface-pack-pipeline-layout"),
213        bind_group_layouts: &[&bind_group_layout],
214        push_constant_ranges: &[],
215    });
216
217    let pipeline =
218        device.create_compute_pipeline(&crate::wgpu_compat::wgpu_compute_pipeline_descriptor! {
219            label: Some("surface-pack-pipeline"),
220            layout: Some(&pipeline_layout),
221            module: &shader,
222            entry_point: "main",
223        });
224
225    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
226        label: Some("surface-pack-bind-group"),
227        layout: &bind_group_layout,
228        entries: &[
229            wgpu::BindGroupEntry {
230                binding: 0,
231                resource: x_buffer.as_ref().as_entire_binding(),
232            },
233            wgpu::BindGroupEntry {
234                binding: 1,
235                resource: y_buffer.as_ref().as_entire_binding(),
236            },
237            wgpu::BindGroupEntry {
238                binding: 2,
239                resource: inputs.z_buffer.as_ref().as_entire_binding(),
240            },
241            wgpu::BindGroupEntry {
242                binding: 3,
243                resource: color_buffer.as_ref().as_entire_binding(),
244            },
245            wgpu::BindGroupEntry {
246                binding: 4,
247                resource: output_buffer.as_ref().as_entire_binding(),
248            },
249            wgpu::BindGroupEntry {
250                binding: 5,
251                resource: uniform_buffer.as_entire_binding(),
252            },
253        ],
254    });
255
256    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
257        label: Some("surface-pack-encoder"),
258    });
259    {
260        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
261            label: Some("surface-pack-pass"),
262            timestamp_writes: None,
263        });
264        pass.set_pipeline(&pipeline);
265        pass.set_bind_group(0, &bind_group, &[]);
266        let workgroups = vertex_count.div_ceil(workgroup_size);
267        pass.dispatch_workgroups(workgroups, 1, 1);
268    }
269    queue.submit(Some(encoder.finish()));
270
271    Ok(GpuVertexBuffer::new(output_buffer, vertex_count as usize))
272}
273
274fn compile_shader(
275    device: &Arc<wgpu::Device>,
276    workgroup_size: u32,
277    scalar: ScalarType,
278) -> wgpu::ShaderModule {
279    let template = match scalar {
280        ScalarType::F32 => shaders::surface::F32,
281        ScalarType::F64 => shaders::surface::F64,
282    };
283    let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
284    device.create_shader_module(wgpu::ShaderModuleDescriptor {
285        label: Some("surface-pack-shader"),
286        source: wgpu::ShaderSource::Wgsl(source.into()),
287    })
288}
289
290#[cfg(test)]
291mod stress_tests {
292    use super::*;
293    use pollster::FutureExt;
294
295    fn maybe_device() -> Option<(Arc<wgpu::Device>, Arc<wgpu::Queue>)> {
296        if std::env::var("RUNMAT_PLOT_SKIP_GPU_TESTS").is_ok()
297            || std::env::var("RUNMAT_PLOT_FORCE_GPU_TESTS").is_err()
298        {
299            return None;
300        }
301        let instance = wgpu::Instance::default();
302        let adapter = instance
303            .request_adapter(&wgpu::RequestAdapterOptions {
304                power_preference: wgpu::PowerPreference::HighPerformance,
305                compatible_surface: None,
306                force_fallback_adapter: false,
307            })
308            .block_on()?;
309        let limits = adapter.limits();
310        let (device, queue) = adapter
311            .request_device(
312                &crate::wgpu_compat::device_descriptor(
313                    Some("runmat-plot-surface-test-device"),
314                    wgpu::Features::empty(),
315                    limits,
316                ),
317                None,
318            )
319            .block_on()
320            .ok()?;
321        Some((Arc::new(device), Arc::new(queue)))
322    }
323
324    #[test]
325    fn gpu_packer_handles_large_surface() {
326        let Some((device, queue)) = maybe_device() else {
327            return;
328        };
329        let x_len = 2048u32;
330        let y_len = 2048u32;
331        let total = (x_len * y_len) as usize;
332        let x_axis: Vec<f32> = (0..x_len).map(|i| i as f32 * 0.1).collect();
333        let y_axis: Vec<f32> = (0..y_len).map(|i| i as f32 * 0.1).collect();
334        let mut z_data = vec![0.0f32; total];
335        for (idx, value) in z_data.iter_mut().enumerate() {
336            let x = (idx % x_len as usize) as f32 * 0.01;
337            let y = (idx / x_len as usize) as f32 * 0.01;
338            *value = (x.sin() + y.cos()) * 0.5;
339        }
340        let z_buffer = Arc::new(
341            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
342                label: Some("surface-test-z"),
343                contents: bytemuck::cast_slice(&z_data),
344                usage: wgpu::BufferUsages::STORAGE,
345            }),
346        );
347
348        let color_table: Vec<[f32; 4]> = (0..256)
349            .map(|i| {
350                let t = i as f32 / 255.0;
351                [t, 1.0 - t, 0.5, 1.0]
352            })
353            .collect();
354
355        let inputs = SurfaceGpuInputs {
356            x_axis: SurfaceAxis::F32(&x_axis),
357            y_axis: SurfaceAxis::F32(&y_axis),
358            z_buffer,
359            color_table: &color_table,
360            x_len,
361            y_len,
362            scalar: ScalarType::F32,
363        };
364        let stride = 8;
365        let lod_x_len = x_len.div_ceil(stride);
366        let lod_y_len = y_len.div_ceil(stride);
367        let params = SurfaceGpuParams {
368            min_z: -1.0,
369            max_z: 1.0,
370            alpha: 1.0,
371            flatten_z: false,
372            x_stride: stride,
373            y_stride: stride,
374            lod_x_len,
375            lod_y_len,
376        };
377
378        let gpu_vertices =
379            pack_surface_vertices(&device, &queue, &inputs, &params).expect("surface pack failed");
380        assert!(gpu_vertices.vertex_count > 0);
381        assert_eq!(gpu_vertices.vertex_count, (lod_x_len * lod_y_len) as usize);
382    }
383}
runmat_plot/gpu/surface.rs

runmat_plot/gpu/
surface.rs
