runmat_plot/gpu/
line.rs

1use crate::core::renderer::Vertex;
2use crate::core::scene::GpuVertexBuffer;
3use crate::gpu::shaders;
4use crate::gpu::{tuning, ScalarType};
5use crate::plots::line::LineStyle;
6use glam::Vec4;
7use log::trace;
8use std::sync::Arc;
9use wgpu::util::DeviceExt;
10
11/// Inputs required to pack line vertices directly on the GPU.
12#[derive(Debug, Clone)]
13pub struct LineGpuInputs {
14    pub x_buffer: Arc<wgpu::Buffer>,
15    pub y_buffer: Arc<wgpu::Buffer>,
16    pub len: u32,
17    pub scalar: ScalarType,
18}
19
20/// Parameters describing how the GPU vertices should be generated.
21pub struct LineGpuParams {
22    pub color: Vec4,
23    /// Half-width in pixels used to extrude thick line triangles in viewport space.
24    pub half_width_px: f32,
25    pub viewport_width_px: f32,
26    pub viewport_height_px: f32,
27    pub x_min: f32,
28    pub x_span: f32,
29    pub y_min: f32,
30    pub y_span: f32,
31    pub line_style: LineStyle,
32    pub marker_size: f32,
33}
34
35#[repr(C)]
36#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
37struct LineSegmentUniforms {
38    color: [f32; 4],
39    count: u32,
40    line_style: u32,
41    half_width_px: f32,
42    _pad0: f32,
43    viewport_width_px: f32,
44    viewport_height_px: f32,
45    x_min: f32,
46    x_span: f32,
47    y_min: f32,
48    y_span: f32,
49    _pad1: [f32; 2],
50}
51
52#[repr(C)]
53#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
54struct MarkerUniforms {
55    color: [f32; 4],
56    count: u32,
57    size: f32,
58    _pad: [u32; 2],
59}
60
61pub fn pack_vertices_from_xy(
62    device: &Arc<wgpu::Device>,
63    queue: &Arc<wgpu::Queue>,
64    inputs: &LineGpuInputs,
65    params: &LineGpuParams,
66) -> Result<GpuVertexBuffer, String> {
67    if inputs.len < 2 {
68        return Err("plot: line inputs must contain at least two points".to_string());
69    }
70
71    let segments = inputs.len - 1;
72    if segments == 0 {
73        return Err("plot: unable to construct segments from degenerate input".to_string());
74    }
75    let vertices_per_segment = 6u64;
76    let max_vertices = segments as u64 * vertices_per_segment;
77    trace!(
78        target: "runmat_plot",
79        "line-pack-kernel: dispatch segments={} max_vertices={} half_width_px={} viewport=({},{})",
80        segments,
81        max_vertices,
82        params.half_width_px,
83        params.viewport_width_px,
84        params.viewport_height_px
85    );
86
87    let workgroup_size = tuning::effective_workgroup_size();
88    let shader = compile_shader(device, workgroup_size, inputs.scalar);
89
90    let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
91        label: Some("line-pack-bind-layout"),
92        entries: &[
93            wgpu::BindGroupLayoutEntry {
94                binding: 0,
95                visibility: wgpu::ShaderStages::COMPUTE,
96                ty: wgpu::BindingType::Buffer {
97                    ty: wgpu::BufferBindingType::Storage { read_only: true },
98                    has_dynamic_offset: false,
99                    min_binding_size: None,
100                },
101                count: None,
102            },
103            wgpu::BindGroupLayoutEntry {
104                binding: 1,
105                visibility: wgpu::ShaderStages::COMPUTE,
106                ty: wgpu::BindingType::Buffer {
107                    ty: wgpu::BufferBindingType::Storage { read_only: true },
108                    has_dynamic_offset: false,
109                    min_binding_size: None,
110                },
111                count: None,
112            },
113            wgpu::BindGroupLayoutEntry {
114                binding: 2,
115                visibility: wgpu::ShaderStages::COMPUTE,
116                ty: wgpu::BindingType::Buffer {
117                    ty: wgpu::BufferBindingType::Storage { read_only: false },
118                    has_dynamic_offset: false,
119                    min_binding_size: None,
120                },
121                count: None,
122            },
123            wgpu::BindGroupLayoutEntry {
124                binding: 3,
125                visibility: wgpu::ShaderStages::COMPUTE,
126                ty: wgpu::BindingType::Buffer {
127                    ty: wgpu::BufferBindingType::Uniform,
128                    has_dynamic_offset: false,
129                    min_binding_size: None,
130                },
131                count: None,
132            },
133        ],
134    });
135
136    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
137        label: Some("line-pack-pipeline-layout"),
138        bind_group_layouts: &[&bind_group_layout],
139        push_constant_ranges: &[],
140    });
141
142    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
143        label: Some("line-pack-pipeline"),
144        layout: Some(&pipeline_layout),
145        module: &shader,
146        entry_point: "main",
147    });
148
149    let output_size = max_vertices * std::mem::size_of::<Vertex>() as u64;
150    let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
151        label: Some("line-gpu-vertices"),
152        size: output_size,
153        usage: wgpu::BufferUsages::STORAGE
154            | wgpu::BufferUsages::VERTEX
155            | wgpu::BufferUsages::COPY_DST,
156        mapped_at_creation: false,
157    }));
158
159    let uniforms = LineSegmentUniforms {
160        color: params.color.to_array(),
161        count: inputs.len,
162        line_style: line_style_code(params.line_style),
163        half_width_px: params.half_width_px.max(0.0),
164        _pad0: 0.0,
165        viewport_width_px: params.viewport_width_px.max(1.0),
166        viewport_height_px: params.viewport_height_px.max(1.0),
167        x_min: params.x_min,
168        x_span: params.x_span.max(1e-12),
169        y_min: params.y_min,
170        y_span: params.y_span.max(1e-12),
171        _pad1: [0.0; 2],
172    };
173    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
174        label: Some("line-pack-uniforms"),
175        contents: bytemuck::bytes_of(&uniforms),
176        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
177    });
178
179    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
180        label: Some("line-pack-bind-group"),
181        layout: &bind_group_layout,
182        entries: &[
183            wgpu::BindGroupEntry {
184                binding: 0,
185                resource: inputs.x_buffer.as_entire_binding(),
186            },
187            wgpu::BindGroupEntry {
188                binding: 1,
189                resource: inputs.y_buffer.as_entire_binding(),
190            },
191            wgpu::BindGroupEntry {
192                binding: 2,
193                resource: output_buffer.as_entire_binding(),
194            },
195            wgpu::BindGroupEntry {
196                binding: 3,
197                resource: uniform_buffer.as_entire_binding(),
198            },
199        ],
200    });
201
202    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
203        label: Some("line-pack-encoder"),
204    });
205    {
206        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
207            label: Some("line-pack-pass"),
208            timestamp_writes: None,
209        });
210        pass.set_pipeline(&pipeline);
211        pass.set_bind_group(0, &bind_group, &[]);
212        let workgroups = segments.div_ceil(workgroup_size);
213        pass.dispatch_workgroups(workgroups, 1, 1);
214    }
215    queue.submit(Some(encoder.finish()));
216    Ok(GpuVertexBuffer::new(output_buffer, max_vertices as usize))
217}
218
219fn compile_shader(
220    device: &Arc<wgpu::Device>,
221    workgroup_size: u32,
222    scalar: ScalarType,
223) -> wgpu::ShaderModule {
224    let template = match scalar {
225        ScalarType::F32 => shaders::line::F32,
226        ScalarType::F64 => shaders::line::F64,
227    };
228    let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
229    device.create_shader_module(wgpu::ShaderModuleDescriptor {
230        label: Some("line-pack-shader"),
231        source: wgpu::ShaderSource::Wgsl(source.into()),
232    })
233}
234
235pub fn pack_marker_vertices_from_xy(
236    device: &Arc<wgpu::Device>,
237    queue: &Arc<wgpu::Queue>,
238    inputs: &LineGpuInputs,
239    params: &LineGpuParams,
240) -> Result<GpuVertexBuffer, String> {
241    if inputs.len < 1 {
242        return Err("plot: marker inputs must contain at least one point".to_string());
243    }
244
245    let workgroup_size = tuning::effective_workgroup_size();
246    let shader = compile_marker_shader(device, workgroup_size, inputs.scalar);
247
248    let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
249        label: Some("line-marker-pack-bind-layout"),
250        entries: &[
251            wgpu::BindGroupLayoutEntry {
252                binding: 0,
253                visibility: wgpu::ShaderStages::COMPUTE,
254                ty: wgpu::BindingType::Buffer {
255                    ty: wgpu::BufferBindingType::Storage { read_only: true },
256                    has_dynamic_offset: false,
257                    min_binding_size: None,
258                },
259                count: None,
260            },
261            wgpu::BindGroupLayoutEntry {
262                binding: 1,
263                visibility: wgpu::ShaderStages::COMPUTE,
264                ty: wgpu::BindingType::Buffer {
265                    ty: wgpu::BufferBindingType::Storage { read_only: true },
266                    has_dynamic_offset: false,
267                    min_binding_size: None,
268                },
269                count: None,
270            },
271            wgpu::BindGroupLayoutEntry {
272                binding: 2,
273                visibility: wgpu::ShaderStages::COMPUTE,
274                ty: wgpu::BindingType::Buffer {
275                    ty: wgpu::BufferBindingType::Storage { read_only: false },
276                    has_dynamic_offset: false,
277                    min_binding_size: None,
278                },
279                count: None,
280            },
281            wgpu::BindGroupLayoutEntry {
282                binding: 3,
283                visibility: wgpu::ShaderStages::COMPUTE,
284                ty: wgpu::BindingType::Buffer {
285                    ty: wgpu::BufferBindingType::Uniform,
286                    has_dynamic_offset: false,
287                    min_binding_size: None,
288                },
289                count: None,
290            },
291        ],
292    });
293
294    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
295        label: Some("line-marker-pack-pipeline-layout"),
296        bind_group_layouts: &[&bind_group_layout],
297        push_constant_ranges: &[],
298    });
299
300    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
301        label: Some("line-marker-pack-pipeline"),
302        layout: Some(&pipeline_layout),
303        module: &shader,
304        entry_point: "main",
305    });
306
307    // Direct point rendering expands each point into a quad (2 triangles = 6 vertices).
308    let expanded_vertices = inputs.len as u64 * 6;
309    let output_size = expanded_vertices * std::mem::size_of::<Vertex>() as u64;
310    let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
311        label: Some("line-marker-gpu-vertices"),
312        size: output_size,
313        usage: wgpu::BufferUsages::STORAGE
314            | wgpu::BufferUsages::VERTEX
315            | wgpu::BufferUsages::COPY_DST,
316        mapped_at_creation: false,
317    }));
318
319    let uniforms = MarkerUniforms {
320        color: params.color.to_array(),
321        count: inputs.len,
322        size: params.marker_size,
323        _pad: [0; 2],
324    };
325    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
326        label: Some("line-marker-pack-uniforms"),
327        contents: bytemuck::bytes_of(&uniforms),
328        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
329    });
330
331    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
332        label: Some("line-marker-pack-bind-group"),
333        layout: &bind_group_layout,
334        entries: &[
335            wgpu::BindGroupEntry {
336                binding: 0,
337                resource: inputs.x_buffer.as_entire_binding(),
338            },
339            wgpu::BindGroupEntry {
340                binding: 1,
341                resource: inputs.y_buffer.as_entire_binding(),
342            },
343            wgpu::BindGroupEntry {
344                binding: 2,
345                resource: output_buffer.as_entire_binding(),
346            },
347            wgpu::BindGroupEntry {
348                binding: 3,
349                resource: uniform_buffer.as_entire_binding(),
350            },
351        ],
352    });
353
354    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
355        label: Some("line-marker-pack-encoder"),
356    });
357    {
358        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
359            label: Some("line-marker-pack-pass"),
360            timestamp_writes: None,
361        });
362        pass.set_pipeline(&pipeline);
363        pass.set_bind_group(0, &bind_group, &[]);
364        let workgroups = inputs.len.div_ceil(workgroup_size);
365        pass.dispatch_workgroups(workgroups, 1, 1);
366    }
367    queue.submit(Some(encoder.finish()));
368
369    Ok(GpuVertexBuffer::new(
370        output_buffer,
371        (inputs.len as usize) * 6,
372    ))
373}
374
375fn compile_marker_shader(
376    device: &Arc<wgpu::Device>,
377    workgroup_size: u32,
378    scalar: ScalarType,
379) -> wgpu::ShaderModule {
380    let template = match scalar {
381        ScalarType::F32 => shaders::line::MARKER_F32,
382        ScalarType::F64 => shaders::line::MARKER_F64,
383    };
384    let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
385    device.create_shader_module(wgpu::ShaderModuleDescriptor {
386        label: Some("line-marker-pack-shader"),
387        source: wgpu::ShaderSource::Wgsl(source.into()),
388    })
389}
390
391fn line_style_code(style: LineStyle) -> u32 {
392    match style {
393        LineStyle::Solid => 0,
394        LineStyle::Dashed => 1,
395        LineStyle::Dotted => 2,
396        LineStyle::DashDot => 3,
397    }
398}
runmat_plot/gpu/line.rs

runmat_plot/gpu/
line.rs
