1use crate::core::renderer::Vertex;
2use crate::core::scene::GpuVertexBuffer;
3use crate::gpu::shaders;
4use crate::gpu::{tuning, ScalarType};
5use crate::plots::line::LineStyle;
6use glam::Vec4;
7use std::sync::Arc;
8use wgpu::util::DeviceExt;
9
10#[derive(Debug, Clone)]
11pub struct Line3GpuInputs {
12 pub x_buffer: Arc<wgpu::Buffer>,
13 pub y_buffer: Arc<wgpu::Buffer>,
14 pub z_buffer: Arc<wgpu::Buffer>,
15 pub len: u32,
16 pub scalar: ScalarType,
17}
18
19pub struct Line3GpuParams {
20 pub color: Vec4,
21 pub half_width_data: f32,
22 pub thick: bool,
23 pub line_style: LineStyle,
24}
25
26#[repr(C)]
27#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
28struct Line3Uniforms {
29 color: [f32; 4],
30 count: u32,
31 half_width_data: f32,
32 line_style: u32,
33 thick: u32,
34 _pad: u32,
35}
36
37pub fn pack_vertices_from_xyz(
38 device: &Arc<wgpu::Device>,
39 queue: &Arc<wgpu::Queue>,
40 inputs: &Line3GpuInputs,
41 params: &Line3GpuParams,
42) -> Result<GpuVertexBuffer, String> {
43 if inputs.len < 2 {
44 return Err("plot3: line inputs must contain at least two points".to_string());
45 }
46 let segments = inputs.len - 1;
47 let vertices_per_segment = if params.thick { 6u64 } else { 2u64 };
48 let max_vertices = segments as u64 * vertices_per_segment;
49 let workgroup_size = tuning::effective_workgroup_size();
50 let shader = compile_shader(device, workgroup_size, inputs.scalar);
51
52 let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
53 label: Some("line3-pack-bind-layout"),
54 entries: &[
55 storage_entry(0, true),
56 storage_entry(1, true),
57 storage_entry(2, true),
58 storage_entry(3, false),
59 uniform_entry(4),
60 ],
61 });
62 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
63 label: Some("line3-pack-pipeline-layout"),
64 bind_group_layouts: &[&bind_group_layout],
65 push_constant_ranges: &[],
66 });
67 let pipeline =
68 device.create_compute_pipeline(&crate::wgpu_compat::wgpu_compute_pipeline_descriptor! {
69 label: Some("line3-pack-pipeline"),
70 layout: Some(&pipeline_layout),
71 module: &shader,
72 entry_point: "main",
73 });
74
75 let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
76 label: Some("line3-gpu-vertices"),
77 size: max_vertices * std::mem::size_of::<Vertex>() as u64,
78 usage: wgpu::BufferUsages::STORAGE
79 | wgpu::BufferUsages::VERTEX
80 | wgpu::BufferUsages::COPY_DST
81 | wgpu::BufferUsages::COPY_SRC,
82 mapped_at_creation: false,
83 }));
84
85 let uniforms = Line3Uniforms {
86 color: params.color.to_array(),
87 count: inputs.len,
88 half_width_data: params.half_width_data.max(0.0001),
89 line_style: line_style_code(params.line_style),
90 thick: if params.thick { 1 } else { 0 },
91 _pad: 0,
92 };
93 let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
94 label: Some("line3-pack-uniforms"),
95 contents: bytemuck::bytes_of(&uniforms),
96 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
97 });
98
99 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
100 label: Some("line3-pack-bind-group"),
101 layout: &bind_group_layout,
102 entries: &[
103 wgpu::BindGroupEntry {
104 binding: 0,
105 resource: inputs.x_buffer.as_entire_binding(),
106 },
107 wgpu::BindGroupEntry {
108 binding: 1,
109 resource: inputs.y_buffer.as_entire_binding(),
110 },
111 wgpu::BindGroupEntry {
112 binding: 2,
113 resource: inputs.z_buffer.as_entire_binding(),
114 },
115 wgpu::BindGroupEntry {
116 binding: 3,
117 resource: output_buffer.as_entire_binding(),
118 },
119 wgpu::BindGroupEntry {
120 binding: 4,
121 resource: uniform_buffer.as_entire_binding(),
122 },
123 ],
124 });
125
126 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
127 label: Some("line3-pack-encoder"),
128 });
129 {
130 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
131 label: Some("line3-pack-pass"),
132 timestamp_writes: None,
133 });
134 pass.set_pipeline(&pipeline);
135 pass.set_bind_group(0, &bind_group, &[]);
136 pass.dispatch_workgroups(segments.div_ceil(workgroup_size), 1, 1);
137 }
138 queue.submit(Some(encoder.finish()));
139 Ok(GpuVertexBuffer::new(output_buffer, max_vertices as usize))
140}
141
142fn compile_shader(
143 device: &Arc<wgpu::Device>,
144 workgroup_size: u32,
145 scalar: ScalarType,
146) -> wgpu::ShaderModule {
147 let template = match scalar {
148 ScalarType::F32 => shaders::line3::F32,
149 ScalarType::F64 => shaders::line3::F64,
150 };
151 let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
152 device.create_shader_module(wgpu::ShaderModuleDescriptor {
153 label: Some("line3-pack-shader"),
154 source: wgpu::ShaderSource::Wgsl(source.into()),
155 })
156}
157
158fn line_style_code(style: LineStyle) -> u32 {
159 match style {
160 LineStyle::Solid => 0,
161 LineStyle::Dashed => 1,
162 LineStyle::Dotted => 2,
163 LineStyle::DashDot => 3,
164 }
165}
166
167fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
168 wgpu::BindGroupLayoutEntry {
169 binding,
170 visibility: wgpu::ShaderStages::COMPUTE,
171 ty: wgpu::BindingType::Buffer {
172 ty: wgpu::BufferBindingType::Storage { read_only },
173 has_dynamic_offset: false,
174 min_binding_size: None,
175 },
176 count: None,
177 }
178}
179
180fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
181 wgpu::BindGroupLayoutEntry {
182 binding,
183 visibility: wgpu::ShaderStages::COMPUTE,
184 ty: wgpu::BindingType::Buffer {
185 ty: wgpu::BufferBindingType::Uniform,
186 has_dynamic_offset: false,
187 min_binding_size: None,
188 },
189 count: None,
190 }
191}