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 =
66        device.create_compute_pipeline(&crate::wgpu_compat::wgpu_compute_pipeline_descriptor! {
67            label: Some("errorbar-pack-pipeline"),
68            layout: Some(&pipeline_layout),
69            module: &shader,
70            entry_point: "main",
71        });
72    let segments_per_point = match params.orientation {
73        0 => 6u64,
74        1 => 6u64,
75        _ => 12u64,
76    };
77    let max_vertices = inputs.len as u64 * segments_per_point;
78    let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
79        label: Some("errorbar-gpu-vertices"),
80        size: max_vertices * std::mem::size_of::<Vertex>() as u64,
81        usage: wgpu::BufferUsages::STORAGE
82            | wgpu::BufferUsages::VERTEX
83            | wgpu::BufferUsages::COPY_DST
84            | wgpu::BufferUsages::COPY_SRC,
85        mapped_at_creation: false,
86    }));
87    let uniforms = ErrorBarUniforms {
88        color: params.color.to_array(),
89        count: inputs.len,
90        line_style: line_style_code(params.line_style),
91        cap_half_width: params.cap_size_data.max(0.0) * 0.5,
92        orientation: params.orientation,
93    };
94    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
95        label: Some("errorbar-pack-uniforms"),
96        contents: bytemuck::bytes_of(&uniforms),
97        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
98    });
99    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
100        label: Some("errorbar-pack-bind-group"),
101        layout: &bind_group_layout,
102        entries: &[
103            wgpu::BindGroupEntry {
104                binding: 0,
105                resource: inputs.x_buffer.as_entire_binding(),
106            },
107            wgpu::BindGroupEntry {
108                binding: 1,
109                resource: inputs.y_buffer.as_entire_binding(),
110            },
111            wgpu::BindGroupEntry {
112                binding: 2,
113                resource: inputs
114                    .x_neg_buffer
115                    .as_ref()
116                    .map(|b| b.as_entire_binding())
117                    .unwrap_or(inputs.x_buffer.as_entire_binding()),
118            },
119            wgpu::BindGroupEntry {
120                binding: 3,
121                resource: inputs
122                    .x_pos_buffer
123                    .as_ref()
124                    .map(|b| b.as_entire_binding())
125                    .unwrap_or(inputs.x_buffer.as_entire_binding()),
126            },
127            wgpu::BindGroupEntry {
128                binding: 4,
129                resource: inputs.y_neg_buffer.as_entire_binding(),
130            },
131            wgpu::BindGroupEntry {
132                binding: 5,
133                resource: inputs.y_pos_buffer.as_entire_binding(),
134            },
135            wgpu::BindGroupEntry {
136                binding: 6,
137                resource: output_buffer.as_entire_binding(),
138            },
139            wgpu::BindGroupEntry {
140                binding: 7,
141                resource: uniform_buffer.as_entire_binding(),
142            },
143        ],
144    });
145    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
146        label: Some("errorbar-pack-encoder"),
147    });
148    {
149        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
150            label: Some("errorbar-pack-pass"),
151            timestamp_writes: None,
152        });
153        pass.set_pipeline(&pipeline);
154        pass.set_bind_group(0, &bind_group, &[]);
155        pass.dispatch_workgroups(inputs.len.div_ceil(workgroup_size), 1, 1);
156    }
157    queue.submit(Some(encoder.finish()));
158    Ok(GpuVertexBuffer::new(output_buffer, max_vertices as usize))
159}
160
161fn compile_shader(
162    device: &Arc<wgpu::Device>,
163    workgroup_size: u32,
164    scalar: ScalarType,
165) -> wgpu::ShaderModule {
166    let template = match scalar {
167        ScalarType::F32 => shaders::errorbar::F32,
168        ScalarType::F64 => shaders::errorbar::F64,
169    };
170    let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
171    device.create_shader_module(wgpu::ShaderModuleDescriptor {
172        label: Some("errorbar-pack-shader"),
173        source: wgpu::ShaderSource::Wgsl(source.into()),
174    })
175}
176
177fn line_style_code(style: LineStyle) -> u32 {
178    match style {
179        LineStyle::Solid => 0,
180        LineStyle::Dashed => 1,
181        LineStyle::Dotted => 2,
182        LineStyle::DashDot => 3,
183    }
184}
185fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
186    wgpu::BindGroupLayoutEntry {
187        binding,
188        visibility: wgpu::ShaderStages::COMPUTE,
189        ty: wgpu::BindingType::Buffer {
190            ty: wgpu::BufferBindingType::Storage { read_only },
191            has_dynamic_offset: false,
192            min_binding_size: None,
193        },
194        count: None,
195    }
196}
197fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
198    wgpu::BindGroupLayoutEntry {
199        binding,
200        visibility: wgpu::ShaderStages::COMPUTE,
201        ty: wgpu::BindingType::Buffer {
202            ty: wgpu::BufferBindingType::Uniform,
203            has_dynamic_offset: false,
204            min_binding_size: None,
205        },
206        count: None,
207    }
208}
runmat_plot/gpu/errorbar.rs

runmat_plot/gpu/
errorbar.rs
