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