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