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

runmat_plot/gpu/
stairs.rs
