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