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