runmat_plot/gpu/
stairs.rs

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

runmat_plot/gpu/
stairs.rs
