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