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