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