1use crate::core::renderer::Vertex;
2use crate::core::scene::GpuVertexBuffer;
3use crate::gpu::axis::{axis_storage_buffer, AxisData};
4use crate::gpu::shaders;
5use crate::gpu::{tuning, ScalarType};
6use glam::Vec4;
7use std::sync::Arc;
8use wgpu::util::DeviceExt;
9
10pub struct QuiverGpuInputs<'a> {
11 pub x_data: AxisData<'a>,
12 pub y_data: AxisData<'a>,
13 pub u_buffer: Arc<wgpu::Buffer>,
14 pub v_buffer: Arc<wgpu::Buffer>,
15 pub count: u32,
16 pub rows: u32,
17 pub cols: u32,
18 pub xy_mode: u32,
19 pub scalar: ScalarType,
20}
21
22pub struct QuiverGpuParams {
23 pub color: Vec4,
24 pub scale: f32,
25 pub head_size: f32,
26}
27
28#[repr(C)]
29#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
30struct QuiverUniforms {
31 color: [f32; 4],
32 count: u32,
33 rows: u32,
34 cols: u32,
35 xy_mode: u32,
36 scale: f32,
37 head_size: f32,
38 _pad: f32,
39}
40
41pub fn pack_vertices(
42 device: &Arc<wgpu::Device>,
43 queue: &Arc<wgpu::Queue>,
44 inputs: &QuiverGpuInputs<'_>,
45 params: &QuiverGpuParams,
46) -> Result<GpuVertexBuffer, String> {
47 let workgroup_size = tuning::effective_workgroup_size();
48 let shader = compile_shader(device, workgroup_size, inputs.scalar);
49 let x_buffer = axis_storage_buffer(device, "quiver-x", &inputs.x_data, inputs.scalar)?;
50 let y_buffer = axis_storage_buffer(device, "quiver-y", &inputs.y_data, inputs.scalar)?;
51 let vertex_count = inputs.count as u64 * 6;
52 let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
53 label: Some("quiver-gpu-vertices"),
54 size: vertex_count * std::mem::size_of::<Vertex>() as u64,
55 usage: wgpu::BufferUsages::STORAGE
56 | wgpu::BufferUsages::VERTEX
57 | wgpu::BufferUsages::COPY_DST
58 | wgpu::BufferUsages::COPY_SRC,
59 mapped_at_creation: false,
60 }));
61 let uniforms = QuiverUniforms {
62 color: params.color.to_array(),
63 count: inputs.count,
64 rows: inputs.rows,
65 cols: inputs.cols,
66 xy_mode: inputs.xy_mode,
67 scale: params.scale,
68 head_size: params.head_size,
69 _pad: 0.0,
70 };
71 let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
72 label: Some("quiver-pack-uniforms"),
73 contents: bytemuck::bytes_of(&uniforms),
74 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
75 });
76 let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
77 label: Some("quiver-pack-bind-layout"),
78 entries: &[
79 storage_entry(0, true),
80 storage_entry(1, true),
81 storage_entry(2, true),
82 storage_entry(3, true),
83 storage_entry(4, false),
84 uniform_entry(5),
85 ],
86 });
87 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
88 label: Some("quiver-pack-pipeline-layout"),
89 bind_group_layouts: &[&bind_group_layout],
90 push_constant_ranges: &[],
91 });
92 let pipeline =
93 device.create_compute_pipeline(&crate::wgpu_compat::wgpu_compute_pipeline_descriptor! {
94 label: Some("quiver-pack-pipeline"),
95 layout: Some(&pipeline_layout),
96 module: &shader,
97 entry_point: "main",
98 });
99 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
100 label: Some("quiver-pack-bind-group"),
101 layout: &bind_group_layout,
102 entries: &[
103 wgpu::BindGroupEntry {
104 binding: 0,
105 resource: x_buffer.as_entire_binding(),
106 },
107 wgpu::BindGroupEntry {
108 binding: 1,
109 resource: y_buffer.as_entire_binding(),
110 },
111 wgpu::BindGroupEntry {
112 binding: 2,
113 resource: inputs.u_buffer.as_entire_binding(),
114 },
115 wgpu::BindGroupEntry {
116 binding: 3,
117 resource: inputs.v_buffer.as_entire_binding(),
118 },
119 wgpu::BindGroupEntry {
120 binding: 4,
121 resource: output_buffer.as_entire_binding(),
122 },
123 wgpu::BindGroupEntry {
124 binding: 5,
125 resource: uniform_buffer.as_entire_binding(),
126 },
127 ],
128 });
129 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
130 label: Some("quiver-pack-encoder"),
131 });
132 {
133 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
134 label: Some("quiver-pack-pass"),
135 timestamp_writes: None,
136 });
137 pass.set_pipeline(&pipeline);
138 pass.set_bind_group(0, &bind_group, &[]);
139 pass.dispatch_workgroups(inputs.count.div_ceil(workgroup_size), 1, 1);
140 }
141 queue.submit(Some(encoder.finish()));
142 Ok(GpuVertexBuffer::new(output_buffer, vertex_count as usize))
143}
144
145fn compile_shader(
146 device: &Arc<wgpu::Device>,
147 workgroup_size: u32,
148 scalar: ScalarType,
149) -> wgpu::ShaderModule {
150 let template = match scalar {
151 ScalarType::F32 => shaders::quiver::F32,
152 ScalarType::F64 => shaders::quiver::F64,
153 };
154 let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
155 device.create_shader_module(wgpu::ShaderModuleDescriptor {
156 label: Some("quiver-pack-shader"),
157 source: wgpu::ShaderSource::Wgsl(source.into()),
158 })
159}
160
161fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
162 wgpu::BindGroupLayoutEntry {
163 binding,
164 visibility: wgpu::ShaderStages::COMPUTE,
165 ty: wgpu::BindingType::Buffer {
166 ty: wgpu::BufferBindingType::Storage { read_only },
167 has_dynamic_offset: false,
168 min_binding_size: None,
169 },
170 count: None,
171 }
172}
173
174fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
175 wgpu::BindGroupLayoutEntry {
176 binding,
177 visibility: wgpu::ShaderStages::COMPUTE,
178 ty: wgpu::BindingType::Buffer {
179 ty: wgpu::BufferBindingType::Uniform,
180 has_dynamic_offset: false,
181 min_binding_size: None,
182 },
183 count: None,
184 }
185}
186
187#[cfg(test)]
188mod tests {
189 use super::*;
190 use pollster::FutureExt;
191 fn maybe_device() -> Option<(Arc<wgpu::Device>, Arc<wgpu::Queue>)> {
192 if std::env::var("RUNMAT_PLOT_SKIP_GPU_TESTS").is_ok()
193 || std::env::var("RUNMAT_PLOT_FORCE_GPU_TESTS").is_err()
194 {
195 return None;
196 }
197 let instance = wgpu::Instance::default();
198 let adapter = instance
199 .request_adapter(&wgpu::RequestAdapterOptions {
200 power_preference: wgpu::PowerPreference::HighPerformance,
201 compatible_surface: None,
202 force_fallback_adapter: false,
203 })
204 .block_on()?;
205 let (device, queue) = adapter
206 .request_device(
207 &crate::wgpu_compat::device_descriptor(
208 Some("runmat-plot-quiver-test-device"),
209 wgpu::Features::empty(),
210 adapter.limits(),
211 ),
212 None,
213 )
214 .block_on()
215 .ok()?;
216 Some((Arc::new(device), Arc::new(queue)))
217 }
218
219 #[test]
220 fn gpu_packer_generates_quiver_vertices() {
221 let Some((device, queue)) = maybe_device() else {
222 return;
223 };
224 let x = [1.0f32, 2.0f32];
225 let y = [1.0f32, 2.0f32];
226 let u = Arc::new(
227 device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
228 label: Some("quiver-test-u"),
229 contents: bytemuck::cast_slice(&[0.5f32, -0.25f32]),
230 usage: wgpu::BufferUsages::STORAGE,
231 }),
232 );
233 let v = Arc::new(
234 device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
235 label: Some("quiver-test-v"),
236 contents: bytemuck::cast_slice(&[1.0f32, 0.75f32]),
237 usage: wgpu::BufferUsages::STORAGE,
238 }),
239 );
240 let packed = pack_vertices(
241 &device,
242 &queue,
243 &QuiverGpuInputs {
244 x_data: AxisData::F32(&x),
245 y_data: AxisData::F32(&y),
246 u_buffer: u,
247 v_buffer: v,
248 count: 2,
249 rows: 2,
250 cols: 1,
251 xy_mode: 0,
252 scalar: ScalarType::F32,
253 },
254 &QuiverGpuParams {
255 color: Vec4::ONE,
256 scale: 1.0,
257 head_size: 0.2,
258 },
259 )
260 .expect("quiver pack should succeed");
261 assert_eq!(packed.vertex_count, 12);
262 }
263}