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

runmat_plot/gpu/
line.rs
