runmat_plot/gpu/
errorbar.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 std::sync::Arc;
8use wgpu::util::DeviceExt;
9
10#[derive(Debug, Clone)]
11pub struct ErrorBarGpuInputs {
12    pub x_buffer: Arc<wgpu::Buffer>,
13    pub y_buffer: Arc<wgpu::Buffer>,
14    pub x_neg_buffer: Option<Arc<wgpu::Buffer>>,
15    pub x_pos_buffer: Option<Arc<wgpu::Buffer>>,
16    pub y_neg_buffer: Arc<wgpu::Buffer>,
17    pub y_pos_buffer: Arc<wgpu::Buffer>,
18    pub len: u32,
19    pub scalar: ScalarType,
20}
21
22pub struct ErrorBarGpuParams {
23    pub color: Vec4,
24    pub cap_size_data: f32,
25    pub line_style: LineStyle,
26    pub orientation: u32,
27}
28
29#[repr(C)]
30#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
31struct ErrorBarUniforms {
32    color: [f32; 4],
33    count: u32,
34    line_style: u32,
35    cap_half_width: f32,
36    orientation: u32,
37}
38
39pub fn pack_vertical_vertices(
40    device: &Arc<wgpu::Device>,
41    queue: &Arc<wgpu::Queue>,
42    inputs: &ErrorBarGpuInputs,
43    params: &ErrorBarGpuParams,
44) -> Result<GpuVertexBuffer, String> {
45    let workgroup_size = tuning::effective_workgroup_size();
46    let shader = compile_shader(device, workgroup_size, inputs.scalar);
47    let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
48        label: Some("errorbar-pack-bind-layout"),
49        entries: &[
50            storage_entry(0, true),
51            storage_entry(1, true),
52            storage_entry(2, true),
53            storage_entry(3, true),
54            storage_entry(4, true),
55            storage_entry(5, true),
56            storage_entry(6, false),
57            uniform_entry(7),
58        ],
59    });
60    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
61        label: Some("errorbar-pack-pipeline-layout"),
62        bind_group_layouts: &[&bind_group_layout],
63        push_constant_ranges: &[],
64    });
65    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
66        label: Some("errorbar-pack-pipeline"),
67        layout: Some(&pipeline_layout),
68        module: &shader,
69        entry_point: "main",
70    });
71    let segments_per_point = match params.orientation {
72        0 => 6u64,
73        1 => 6u64,
74        _ => 12u64,
75    };
76    let max_vertices = inputs.len as u64 * segments_per_point;
77    let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
78        label: Some("errorbar-gpu-vertices"),
79        size: max_vertices * std::mem::size_of::<Vertex>() as u64,
80        usage: wgpu::BufferUsages::STORAGE
81            | wgpu::BufferUsages::VERTEX
82            | wgpu::BufferUsages::COPY_DST,
83        mapped_at_creation: false,
84    }));
85    let uniforms = ErrorBarUniforms {
86        color: params.color.to_array(),
87        count: inputs.len,
88        line_style: line_style_code(params.line_style),
89        cap_half_width: params.cap_size_data.max(0.0) * 0.5,
90        orientation: params.orientation,
91    };
92    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
93        label: Some("errorbar-pack-uniforms"),
94        contents: bytemuck::bytes_of(&uniforms),
95        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
96    });
97    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
98        label: Some("errorbar-pack-bind-group"),
99        layout: &bind_group_layout,
100        entries: &[
101            wgpu::BindGroupEntry {
102                binding: 0,
103                resource: inputs.x_buffer.as_entire_binding(),
104            },
105            wgpu::BindGroupEntry {
106                binding: 1,
107                resource: inputs.y_buffer.as_entire_binding(),
108            },
109            wgpu::BindGroupEntry {
110                binding: 2,
111                resource: inputs
112                    .x_neg_buffer
113                    .as_ref()
114                    .map(|b| b.as_entire_binding())
115                    .unwrap_or(inputs.x_buffer.as_entire_binding()),
116            },
117            wgpu::BindGroupEntry {
118                binding: 3,
119                resource: inputs
120                    .x_pos_buffer
121                    .as_ref()
122                    .map(|b| b.as_entire_binding())
123                    .unwrap_or(inputs.x_buffer.as_entire_binding()),
124            },
125            wgpu::BindGroupEntry {
126                binding: 4,
127                resource: inputs.y_neg_buffer.as_entire_binding(),
128            },
129            wgpu::BindGroupEntry {
130                binding: 5,
131                resource: inputs.y_pos_buffer.as_entire_binding(),
132            },
133            wgpu::BindGroupEntry {
134                binding: 6,
135                resource: output_buffer.as_entire_binding(),
136            },
137            wgpu::BindGroupEntry {
138                binding: 7,
139                resource: uniform_buffer.as_entire_binding(),
140            },
141        ],
142    });
143    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
144        label: Some("errorbar-pack-encoder"),
145    });
146    {
147        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
148            label: Some("errorbar-pack-pass"),
149            timestamp_writes: None,
150        });
151        pass.set_pipeline(&pipeline);
152        pass.set_bind_group(0, &bind_group, &[]);
153        pass.dispatch_workgroups(inputs.len.div_ceil(workgroup_size), 1, 1);
154    }
155    queue.submit(Some(encoder.finish()));
156    Ok(GpuVertexBuffer::new(output_buffer, max_vertices as usize))
157}
158
159fn compile_shader(
160    device: &Arc<wgpu::Device>,
161    workgroup_size: u32,
162    scalar: ScalarType,
163) -> wgpu::ShaderModule {
164    let template = match scalar {
165        ScalarType::F32 => shaders::errorbar::F32,
166        ScalarType::F64 => shaders::errorbar::F64,
167    };
168    let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
169    device.create_shader_module(wgpu::ShaderModuleDescriptor {
170        label: Some("errorbar-pack-shader"),
171        source: wgpu::ShaderSource::Wgsl(source.into()),
172    })
173}
174
175fn line_style_code(style: LineStyle) -> u32 {
176    match style {
177        LineStyle::Solid => 0,
178        LineStyle::Dashed => 1,
179        LineStyle::Dotted => 2,
180        LineStyle::DashDot => 3,
181    }
182}
183fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
184    wgpu::BindGroupLayoutEntry {
185        binding,
186        visibility: wgpu::ShaderStages::COMPUTE,
187        ty: wgpu::BindingType::Buffer {
188            ty: wgpu::BufferBindingType::Storage { read_only },
189            has_dynamic_offset: false,
190            min_binding_size: None,
191        },
192        count: None,
193    }
194}
195fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
196    wgpu::BindGroupLayoutEntry {
197        binding,
198        visibility: wgpu::ShaderStages::COMPUTE,
199        ty: wgpu::BindingType::Buffer {
200            ty: wgpu::BufferBindingType::Uniform,
201            has_dynamic_offset: false,
202            min_binding_size: None,
203        },
204        count: None,
205    }
206}
runmat_plot/gpu/errorbar.rs

runmat_plot/gpu/
errorbar.rs
