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 =
143        device.create_compute_pipeline(&crate::wgpu_compat::wgpu_compute_pipeline_descriptor! {
144            label: Some("line-pack-pipeline"),
145            layout: Some(&pipeline_layout),
146            module: &shader,
147            entry_point: "main",
148        });
149
150    let output_size = max_vertices * std::mem::size_of::<Vertex>() as u64;
151    let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
152        label: Some("line-gpu-vertices"),
153        size: output_size,
154        usage: wgpu::BufferUsages::STORAGE
155            | wgpu::BufferUsages::VERTEX
156            | wgpu::BufferUsages::COPY_DST
157            | wgpu::BufferUsages::COPY_SRC,
158        mapped_at_creation: false,
159    }));
160
161    let uniforms = LineSegmentUniforms {
162        color: params.color.to_array(),
163        count: inputs.len,
164        line_style: line_style_code(params.line_style),
165        half_width_px: params.half_width_px.max(0.0),
166        _pad0: 0.0,
167        viewport_width_px: params.viewport_width_px.max(1.0),
168        viewport_height_px: params.viewport_height_px.max(1.0),
169        x_min: params.x_min,
170        x_span: params.x_span.max(1e-12),
171        y_min: params.y_min,
172        y_span: params.y_span.max(1e-12),
173        _pad1: [0.0; 2],
174    };
175    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
176        label: Some("line-pack-uniforms"),
177        contents: bytemuck::bytes_of(&uniforms),
178        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
179    });
180
181    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
182        label: Some("line-pack-bind-group"),
183        layout: &bind_group_layout,
184        entries: &[
185            wgpu::BindGroupEntry {
186                binding: 0,
187                resource: inputs.x_buffer.as_entire_binding(),
188            },
189            wgpu::BindGroupEntry {
190                binding: 1,
191                resource: inputs.y_buffer.as_entire_binding(),
192            },
193            wgpu::BindGroupEntry {
194                binding: 2,
195                resource: output_buffer.as_entire_binding(),
196            },
197            wgpu::BindGroupEntry {
198                binding: 3,
199                resource: uniform_buffer.as_entire_binding(),
200            },
201        ],
202    });
203
204    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
205        label: Some("line-pack-encoder"),
206    });
207    {
208        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
209            label: Some("line-pack-pass"),
210            timestamp_writes: None,
211        });
212        pass.set_pipeline(&pipeline);
213        pass.set_bind_group(0, &bind_group, &[]);
214        let workgroups = segments.div_ceil(workgroup_size);
215        pass.dispatch_workgroups(workgroups, 1, 1);
216    }
217    queue.submit(Some(encoder.finish()));
218    Ok(GpuVertexBuffer::new(output_buffer, max_vertices as usize))
219}
220
221fn compile_shader(
222    device: &Arc<wgpu::Device>,
223    workgroup_size: u32,
224    scalar: ScalarType,
225) -> wgpu::ShaderModule {
226    let template = match scalar {
227        ScalarType::F32 => shaders::line::F32,
228        ScalarType::F64 => shaders::line::F64,
229    };
230    let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
231    device.create_shader_module(wgpu::ShaderModuleDescriptor {
232        label: Some("line-pack-shader"),
233        source: wgpu::ShaderSource::Wgsl(source.into()),
234    })
235}
236
237pub fn pack_marker_vertices_from_xy(
238    device: &Arc<wgpu::Device>,
239    queue: &Arc<wgpu::Queue>,
240    inputs: &LineGpuInputs,
241    params: &LineGpuParams,
242) -> Result<GpuVertexBuffer, String> {
243    if inputs.len < 1 {
244        return Err("plot: marker inputs must contain at least one point".to_string());
245    }
246
247    let workgroup_size = tuning::effective_workgroup_size();
248    let shader = compile_marker_shader(device, workgroup_size, inputs.scalar);
249
250    let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
251        label: Some("line-marker-pack-bind-layout"),
252        entries: &[
253            wgpu::BindGroupLayoutEntry {
254                binding: 0,
255                visibility: wgpu::ShaderStages::COMPUTE,
256                ty: wgpu::BindingType::Buffer {
257                    ty: wgpu::BufferBindingType::Storage { read_only: true },
258                    has_dynamic_offset: false,
259                    min_binding_size: None,
260                },
261                count: None,
262            },
263            wgpu::BindGroupLayoutEntry {
264                binding: 1,
265                visibility: wgpu::ShaderStages::COMPUTE,
266                ty: wgpu::BindingType::Buffer {
267                    ty: wgpu::BufferBindingType::Storage { read_only: true },
268                    has_dynamic_offset: false,
269                    min_binding_size: None,
270                },
271                count: None,
272            },
273            wgpu::BindGroupLayoutEntry {
274                binding: 2,
275                visibility: wgpu::ShaderStages::COMPUTE,
276                ty: wgpu::BindingType::Buffer {
277                    ty: wgpu::BufferBindingType::Storage { read_only: false },
278                    has_dynamic_offset: false,
279                    min_binding_size: None,
280                },
281                count: None,
282            },
283            wgpu::BindGroupLayoutEntry {
284                binding: 3,
285                visibility: wgpu::ShaderStages::COMPUTE,
286                ty: wgpu::BindingType::Buffer {
287                    ty: wgpu::BufferBindingType::Uniform,
288                    has_dynamic_offset: false,
289                    min_binding_size: None,
290                },
291                count: None,
292            },
293        ],
294    });
295
296    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
297        label: Some("line-marker-pack-pipeline-layout"),
298        bind_group_layouts: &[&bind_group_layout],
299        push_constant_ranges: &[],
300    });
301
302    let pipeline =
303        device.create_compute_pipeline(&crate::wgpu_compat::wgpu_compute_pipeline_descriptor! {
304            label: Some("line-marker-pack-pipeline"),
305            layout: Some(&pipeline_layout),
306            module: &shader,
307            entry_point: "main",
308        });
309
310    // Direct point rendering expands each point into a quad (2 triangles = 6 vertices).
311    let expanded_vertices = inputs.len as u64 * 6;
312    let output_size = expanded_vertices * std::mem::size_of::<Vertex>() as u64;
313    let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
314        label: Some("line-marker-gpu-vertices"),
315        size: output_size,
316        usage: wgpu::BufferUsages::STORAGE
317            | wgpu::BufferUsages::VERTEX
318            | wgpu::BufferUsages::COPY_DST
319            | wgpu::BufferUsages::COPY_SRC,
320        mapped_at_creation: false,
321    }));
322
323    let uniforms = MarkerUniforms {
324        color: params.color.to_array(),
325        count: inputs.len,
326        size: params.marker_size,
327        _pad: [0; 2],
328    };
329    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
330        label: Some("line-marker-pack-uniforms"),
331        contents: bytemuck::bytes_of(&uniforms),
332        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
333    });
334
335    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
336        label: Some("line-marker-pack-bind-group"),
337        layout: &bind_group_layout,
338        entries: &[
339            wgpu::BindGroupEntry {
340                binding: 0,
341                resource: inputs.x_buffer.as_entire_binding(),
342            },
343            wgpu::BindGroupEntry {
344                binding: 1,
345                resource: inputs.y_buffer.as_entire_binding(),
346            },
347            wgpu::BindGroupEntry {
348                binding: 2,
349                resource: output_buffer.as_entire_binding(),
350            },
351            wgpu::BindGroupEntry {
352                binding: 3,
353                resource: uniform_buffer.as_entire_binding(),
354            },
355        ],
356    });
357
358    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
359        label: Some("line-marker-pack-encoder"),
360    });
361    {
362        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
363            label: Some("line-marker-pack-pass"),
364            timestamp_writes: None,
365        });
366        pass.set_pipeline(&pipeline);
367        pass.set_bind_group(0, &bind_group, &[]);
368        let workgroups = inputs.len.div_ceil(workgroup_size);
369        pass.dispatch_workgroups(workgroups, 1, 1);
370    }
371    queue.submit(Some(encoder.finish()));
372
373    Ok(GpuVertexBuffer::new(
374        output_buffer,
375        (inputs.len as usize) * 6,
376    ))
377}
378
379fn compile_marker_shader(
380    device: &Arc<wgpu::Device>,
381    workgroup_size: u32,
382    scalar: ScalarType,
383) -> wgpu::ShaderModule {
384    let template = match scalar {
385        ScalarType::F32 => shaders::line::MARKER_F32,
386        ScalarType::F64 => shaders::line::MARKER_F64,
387    };
388    let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
389    device.create_shader_module(wgpu::ShaderModuleDescriptor {
390        label: Some("line-marker-pack-shader"),
391        source: wgpu::ShaderSource::Wgsl(source.into()),
392    })
393}
394
395fn line_style_code(style: LineStyle) -> u32 {
396    match style {
397        LineStyle::Solid => 0,
398        LineStyle::Dashed => 1,
399        LineStyle::Dotted => 2,
400        LineStyle::DashDot => 3,
401    }
402}
runmat_plot/gpu/line.rs

runmat_plot/gpu/
line.rs
