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 = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
68 label: Some("line3-pack-pipeline"),
69 layout: Some(&pipeline_layout),
70 module: &shader,
71 entry_point: "main",
72 });
73
74 let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
75 label: Some("line3-gpu-vertices"),
76 size: max_vertices * std::mem::size_of::<Vertex>() as u64,
77 usage: wgpu::BufferUsages::STORAGE
78 | wgpu::BufferUsages::VERTEX
79 | wgpu::BufferUsages::COPY_DST,
80 mapped_at_creation: false,
81 }));
82
83 let uniforms = Line3Uniforms {
84 color: params.color.to_array(),
85 count: inputs.len,
86 half_width_data: params.half_width_data.max(0.0001),
87 line_style: line_style_code(params.line_style),
88 thick: if params.thick { 1 } else { 0 },
89 _pad: 0,
90 };
91 let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
92 label: Some("line3-pack-uniforms"),
93 contents: bytemuck::bytes_of(&uniforms),
94 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
95 });
96
97 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
98 label: Some("line3-pack-bind-group"),
99 layout: &bind_group_layout,
100 entries: &[
101 wgpu::BindGroupEntry {
102 binding: 0,
103 resource: inputs.x_buffer.as_entire_binding(),
104 },
105 wgpu::BindGroupEntry {
106 binding: 1,
107 resource: inputs.y_buffer.as_entire_binding(),
108 },
109 wgpu::BindGroupEntry {
110 binding: 2,
111 resource: inputs.z_buffer.as_entire_binding(),
112 },
113 wgpu::BindGroupEntry {
114 binding: 3,
115 resource: output_buffer.as_entire_binding(),
116 },
117 wgpu::BindGroupEntry {
118 binding: 4,
119 resource: uniform_buffer.as_entire_binding(),
120 },
121 ],
122 });
123
124 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
125 label: Some("line3-pack-encoder"),
126 });
127 {
128 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
129 label: Some("line3-pack-pass"),
130 timestamp_writes: None,
131 });
132 pass.set_pipeline(&pipeline);
133 pass.set_bind_group(0, &bind_group, &[]);
134 pass.dispatch_workgroups(segments.div_ceil(workgroup_size), 1, 1);
135 }
136 queue.submit(Some(encoder.finish()));
137 Ok(GpuVertexBuffer::new(output_buffer, max_vertices as usize))
138}
139
140fn compile_shader(
141 device: &Arc<wgpu::Device>,
142 workgroup_size: u32,
143 scalar: ScalarType,
144) -> wgpu::ShaderModule {
145 let template = match scalar {
146 ScalarType::F32 => shaders::line3::F32,
147 ScalarType::F64 => shaders::line3::F64,
148 };
149 let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
150 device.create_shader_module(wgpu::ShaderModuleDescriptor {
151 label: Some("line3-pack-shader"),
152 source: wgpu::ShaderSource::Wgsl(source.into()),
153 })
154}
155
156fn line_style_code(style: LineStyle) -> u32 {
157 match style {
158 LineStyle::Solid => 0,
159 LineStyle::Dashed => 1,
160 LineStyle::Dotted => 2,
161 LineStyle::DashDot => 3,
162 }
163}
164
165fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
166 wgpu::BindGroupLayoutEntry {
167 binding,
168 visibility: wgpu::ShaderStages::COMPUTE,
169 ty: wgpu::BindingType::Buffer {
170 ty: wgpu::BufferBindingType::Storage { read_only },
171 has_dynamic_offset: false,
172 min_binding_size: None,
173 },
174 count: None,
175 }
176}
177
178fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
179 wgpu::BindGroupLayoutEntry {
180 binding,
181 visibility: wgpu::ShaderStages::COMPUTE,
182 ty: wgpu::BindingType::Buffer {
183 ty: wgpu::BufferBindingType::Uniform,
184 has_dynamic_offset: false,
185 min_binding_size: None,
186 },
187 count: None,
188 }
189}