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 = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
85 label: Some("area-pack-pipeline"),
86 layout: Some(&pipeline_layout),
87 module: &shader,
88 entry_point: "main",
89 });
90 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
91 label: Some("area-pack-bind-group"),
92 layout: &bind_group_layout,
93 entries: &[
94 wgpu::BindGroupEntry {
95 binding: 0,
96 resource: x_buffer.as_entire_binding(),
97 },
98 wgpu::BindGroupEntry {
99 binding: 1,
100 resource: inputs.y_buffer.as_entire_binding(),
101 },
102 wgpu::BindGroupEntry {
103 binding: 2,
104 resource: output_buffer.as_entire_binding(),
105 },
106 wgpu::BindGroupEntry {
107 binding: 3,
108 resource: uniform_buffer.as_entire_binding(),
109 },
110 ],
111 });
112 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
113 label: Some("area-pack-encoder"),
114 });
115 {
116 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
117 label: Some("area-pack-pass"),
118 timestamp_writes: None,
119 });
120 pass.set_pipeline(&pipeline);
121 pass.set_bind_group(0, &bind_group, &[]);
122 pass.dispatch_workgroups((inputs.rows - 1).div_ceil(workgroup_size), 1, 1);
123 }
124 queue.submit(Some(encoder.finish()));
125 Ok(GpuVertexBuffer::new(output_buffer, vertex_count as usize))
126}
127
128fn compile_shader(
129 device: &Arc<wgpu::Device>,
130 workgroup_size: u32,
131 scalar: ScalarType,
132) -> wgpu::ShaderModule {
133 let template = match scalar {
134 ScalarType::F32 => shaders::area::F32,
135 ScalarType::F64 => shaders::area::F64,
136 };
137 let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
138 device.create_shader_module(wgpu::ShaderModuleDescriptor {
139 label: Some("area-pack-shader"),
140 source: wgpu::ShaderSource::Wgsl(source.into()),
141 })
142}
143
144fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
145 wgpu::BindGroupLayoutEntry {
146 binding,
147 visibility: wgpu::ShaderStages::COMPUTE,
148 ty: wgpu::BindingType::Buffer {
149 ty: wgpu::BufferBindingType::Storage { read_only },
150 has_dynamic_offset: false,
151 min_binding_size: None,
152 },
153 count: None,
154 }
155}
156
157fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
158 wgpu::BindGroupLayoutEntry {
159 binding,
160 visibility: wgpu::ShaderStages::COMPUTE,
161 ty: wgpu::BindingType::Buffer {
162 ty: wgpu::BufferBindingType::Uniform,
163 has_dynamic_offset: false,
164 min_binding_size: None,
165 },
166 count: None,
167 }
168}
169
170#[cfg(test)]
171mod tests {
172 use super::*;
173 use pollster::FutureExt;
174 fn maybe_device() -> Option<(Arc<wgpu::Device>, Arc<wgpu::Queue>)> {
175 if std::env::var("RUNMAT_PLOT_SKIP_GPU_TESTS").is_ok()
176 || std::env::var("RUNMAT_PLOT_FORCE_GPU_TESTS").is_err()
177 {
178 return None;
179 }
180 let instance = wgpu::Instance::default();
181 let adapter = instance
182 .request_adapter(&wgpu::RequestAdapterOptions {
183 power_preference: wgpu::PowerPreference::HighPerformance,
184 compatible_surface: None,
185 force_fallback_adapter: false,
186 })
187 .block_on()?;
188 let (device, queue) = adapter
189 .request_device(
190 &wgpu::DeviceDescriptor {
191 label: Some("runmat-plot-area-test-device"),
192 required_features: wgpu::Features::empty(),
193 required_limits: adapter.limits(),
194 },
195 None,
196 )
197 .block_on()
198 .ok()?;
199 Some((Arc::new(device), Arc::new(queue)))
200 }
201
202 #[test]
203 fn gpu_packer_generates_area_vertices() {
204 let Some((device, queue)) = maybe_device() else {
205 return;
206 };
207 let x = [1.0f32, 2.0f32, 3.0f32];
208 let y = Arc::new(
209 device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
210 label: Some("area-test-y"),
211 contents: bytemuck::cast_slice(&[1.0f32, 2.0f32, 3.0f32, 0.5f32, 0.5f32, 0.5f32]),
212 usage: wgpu::BufferUsages::STORAGE,
213 }),
214 );
215 let packed = pack_vertices(
216 &device,
217 &queue,
218 &AreaGpuInputs {
219 x_axis: AxisData::F32(&x),
220 y_buffer: y,
221 rows: 3,
222 cols: 2,
223 target_col: 1,
224 scalar: ScalarType::F32,
225 },
226 &AreaGpuParams {
227 color: Vec4::ONE,
228 baseline: 0.0,
229 },
230 )
231 .expect("area pack should succeed");
232 assert_eq!(packed.vertex_count, 12);
233 }
234}