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

runmat_plot/gpu/
line.rs
